

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.dataset.engine.datasets &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.dataset.engine.datasets</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.dataset.engine.datasets</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2019-2022 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ==============================================================================</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">1. This file is an abstraction of the dataset loading class. It contains</span>
<span class="sd">some basic dataset operations(skip, filter, map, batch, ...).</span>
<span class="sd">2. Specific dataset loading classes can be found in datasets_vision.py, datasets_text.py,</span>
<span class="sd">datasets_audio.py, datasets_standard_format.py and dataets_user_defined.py files.</span>
<span class="sd">    datasets_vision.py: contains vision dataset loading classes.</span>
<span class="sd">    datasets_text.py: contains text dataset loading classes.</span>
<span class="sd">    datasets_audio.py: contains audio dataset loading classes.</span>
<span class="sd">    datasets_standard_format.py: contains standard format loading classes which</span>
<span class="sd">                                 any other kinds of datasets can be converted to.</span>
<span class="sd">    dataets_user_defined.py: contains basic classes that help users to define</span>
<span class="sd">                             flexible ways to load dataset.</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">atexit</span>
<span class="kn">import</span> <span class="nn">glob</span>
<span class="kn">import</span> <span class="nn">json</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">signal</span>
<span class="kn">import</span> <span class="nn">stat</span>
<span class="kn">import</span> <span class="nn">time</span>
<span class="kn">import</span> <span class="nn">uuid</span>
<span class="kn">import</span> <span class="nn">multiprocessing</span>
<span class="kn">from</span> <span class="nn">multiprocessing.pool</span> <span class="kn">import</span> <span class="n">RUN</span><span class="p">,</span> <span class="n">TERMINATE</span>
<span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="kn">from</span> <span class="nn">importlib</span> <span class="kn">import</span> <span class="n">import_module</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">import</span> <span class="nn">threading</span>

<span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">import</span> <span class="nn">weakref</span>
<span class="kn">import</span> <span class="nn">platform</span>
<span class="kn">import</span> <span class="nn">psutil</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>

<span class="kn">import</span> <span class="nn">mindspore._c_dataengine</span> <span class="k">as</span> <span class="nn">cde</span>
<span class="kn">from</span> <span class="nn">mindspore._c_expression</span> <span class="kn">import</span> <span class="n">typing</span>

<span class="kn">from</span> <span class="nn">mindspore</span> <span class="kn">import</span> <span class="n">log</span> <span class="k">as</span> <span class="n">logger</span>
<span class="kn">from</span> <span class="nn">mindspore.parallel._ps_context</span> <span class="kn">import</span> <span class="n">_is_role_pserver</span><span class="p">,</span> <span class="n">_is_role_sched</span>
<span class="kn">from</span> <span class="nn">mindspore.dataset.engine.offload</span> <span class="kn">import</span> <span class="n">GetOffloadModel</span>

<span class="kn">import</span> <span class="nn">mindspore.dataset.transforms.py_transforms</span> <span class="k">as</span> <span class="nn">py_transforms</span>
<span class="kn">from</span> <span class="nn">mindspore.dataset.text.utils</span> <span class="kn">import</span> <span class="n">SentencePieceModel</span><span class="p">,</span> <span class="n">DE_C_INTER_SENTENCEPIECE_MODE</span>
<span class="kn">from</span> <span class="nn">mindspore.parallel._utils</span> <span class="kn">import</span> <span class="n">_get_device_num</span>

<span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">samplers</span>
<span class="kn">from</span> <span class="nn">.iterators</span> <span class="kn">import</span> <span class="n">DictIterator</span><span class="p">,</span> <span class="n">TupleIterator</span><span class="p">,</span> <span class="n">DummyIterator</span><span class="p">,</span> <span class="n">check_iterator_cleanup</span><span class="p">,</span> <span class="n">_set_iterator_cleanup</span><span class="p">,</span> \
    <span class="n">ITERATORS_LIST</span><span class="p">,</span> <span class="n">_unset_iterator_cleanup</span>
<span class="kn">from</span> <span class="nn">.queue</span> <span class="kn">import</span> <span class="n">_SharedQueue</span>
<span class="kn">from</span> <span class="nn">.validators</span> <span class="kn">import</span> <span class="n">check_batch</span><span class="p">,</span> <span class="n">check_shuffle</span><span class="p">,</span> <span class="n">check_map</span><span class="p">,</span> <span class="n">check_filter</span><span class="p">,</span> <span class="n">check_repeat</span><span class="p">,</span> <span class="n">check_skip</span><span class="p">,</span> <span class="n">check_zip</span><span class="p">,</span> \
    <span class="n">check_rename</span><span class="p">,</span> <span class="n">check_device_send</span><span class="p">,</span> <span class="n">check_take</span><span class="p">,</span> <span class="n">check_project</span><span class="p">,</span> \
    <span class="n">check_sync_wait</span><span class="p">,</span> <span class="n">check_zip_dataset</span><span class="p">,</span> <span class="n">check_add_column</span><span class="p">,</span> <span class="n">check_concat</span><span class="p">,</span> <span class="n">check_split</span><span class="p">,</span> <span class="n">check_bucket_batch_by_length</span><span class="p">,</span> \
    <span class="n">check_save</span><span class="p">,</span> <span class="n">check_tuple_iterator</span><span class="p">,</span> <span class="n">check_dict_iterator</span><span class="p">,</span> <span class="n">check_schema</span><span class="p">,</span> <span class="n">check_to_device_send</span>
<span class="kn">from</span> <span class="nn">..core.config</span> <span class="kn">import</span> <span class="n">get_callback_timeout</span><span class="p">,</span> <span class="n">_init_device_info</span><span class="p">,</span> <span class="n">get_enable_shared_mem</span><span class="p">,</span> <span class="n">get_num_parallel_workers</span>
<span class="kn">from</span> <span class="nn">..core.datatypes</span> <span class="kn">import</span> <span class="n">mstype_to_detype</span>
<span class="kn">from</span> <span class="nn">..core.validator_helpers</span> <span class="kn">import</span> <span class="n">replace_none</span>
<span class="kn">from</span> <span class="nn">..core.py_util_helpers</span> <span class="kn">import</span> <span class="n">ExceptionHandler</span>
<span class="kn">from</span> <span class="nn">..transforms.py_transforms_util</span> <span class="kn">import</span> <span class="n">FuncWrapper</span>

<span class="k">try</span><span class="p">:</span>
    <span class="n">context</span> <span class="o">=</span> <span class="n">import_module</span><span class="p">(</span><span class="s2">&quot;mindspore.context&quot;</span><span class="p">)</span>
<span class="k">except</span> <span class="ne">ModuleNotFoundError</span><span class="p">:</span>
    <span class="n">context</span> <span class="o">=</span> <span class="kc">None</span>

<span class="k">if</span> <span class="n">platform</span><span class="o">.</span><span class="n">system</span><span class="p">()</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="o">==</span> <span class="s2">&quot;darwin&quot;</span> <span class="ow">and</span> <span class="n">multiprocessing</span><span class="o">.</span><span class="n">get_start_method</span><span class="p">()</span> <span class="o">!=</span> <span class="s2">&quot;fork&quot;</span><span class="p">:</span>
    <span class="n">multiprocessing</span><span class="o">.</span><span class="n">set_start_method</span><span class="p">(</span><span class="s2">&quot;fork&quot;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>


<span class="n">OffloadToManualOffloadMode</span> <span class="o">=</span> <span class="p">{</span>
    <span class="kc">None</span><span class="p">:</span> <span class="n">cde</span><span class="o">.</span><span class="n">ManualOffloadMode</span><span class="o">.</span><span class="n">UNSPECIFIED</span><span class="p">,</span>
    <span class="kc">False</span><span class="p">:</span> <span class="n">cde</span><span class="o">.</span><span class="n">ManualOffloadMode</span><span class="o">.</span><span class="n">DISABLED</span><span class="p">,</span>
    <span class="kc">True</span><span class="p">:</span> <span class="n">cde</span><span class="o">.</span><span class="n">ManualOffloadMode</span><span class="o">.</span><span class="n">ENABLED</span>
<span class="p">}</span>


<span class="k">class</span> <span class="nc">Shuffle</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Specify the shuffle mode.</span>

<span class="sd">    - GLOBAL: Shuffle both the files and samples.</span>
<span class="sd">    - FILES: Shuffle files only.</span>
<span class="sd">    - INFILE: Shuffle data within each file.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">GLOBAL</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;global&quot;</span>
    <span class="n">FILES</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;files&quot;</span>
    <span class="n">INFILE</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;infile&quot;</span>


<span class="n">ShuffleToShuffleMode</span> <span class="o">=</span> <span class="p">{</span><span class="n">Shuffle</span><span class="o">.</span><span class="n">FILES</span><span class="p">:</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleMode</span><span class="o">.</span><span class="n">FILES</span><span class="p">,</span>
                        <span class="n">Shuffle</span><span class="o">.</span><span class="n">GLOBAL</span><span class="p">:</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleMode</span><span class="o">.</span><span class="n">GLOBAL</span><span class="p">,</span>
                        <span class="n">Shuffle</span><span class="o">.</span><span class="n">INFILE</span><span class="p">:</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleMode</span><span class="o">.</span><span class="n">INFILE</span><span class="p">}</span>


<span class="k">def</span> <span class="nf">shuffle_to_shuffle_mode</span><span class="p">(</span><span class="n">shuffle</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Shuffle Enum to Shuffle Mode</span>

<span class="sd">    Args:</span>
<span class="sd">        shuffle (Shuffle): shuffle flag to shuffle mode in C layer</span>

<span class="sd">    Returns:</span>
<span class="sd">        ShuffleMode, shuffle mode</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">shuffle_mode</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleMode</span><span class="o">.</span><span class="n">GLOBAL</span>  <span class="c1"># Global shuffle</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shuffle</span><span class="p">,</span> <span class="n">Shuffle</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">shuffle</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">shuffle</span><span class="p">:</span>
            <span class="n">shuffle_mode</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleMode</span><span class="o">.</span><span class="n">GLOBAL</span>  <span class="c1"># Global shuffle</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">shuffle_mode</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleMode</span><span class="o">.</span><span class="n">FALSE</span>  <span class="c1"># No shuffle</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">shuffle_mode</span> <span class="o">=</span> <span class="n">ShuffleToShuffleMode</span><span class="p">[</span><span class="n">shuffle</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">shuffle_mode</span>


<span class="k">def</span> <span class="nf">shuffle_to_bool</span><span class="p">(</span><span class="n">shuffle</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Shuffle Enum to bool</span>

<span class="sd">    Args:</span>
<span class="sd">        shuffle (Shuffle): shuffle flag to bool</span>

<span class="sd">    Returns:</span>
<span class="sd">        bool, True / False</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">shuffle_bool</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shuffle</span><span class="p">,</span> <span class="n">Shuffle</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">shuffle</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">shuffle_bool</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">elif</span> <span class="n">shuffle</span><span class="p">:</span>
            <span class="n">shuffle_bool</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">shuffle_bool</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">shuffle_bool</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">return</span> <span class="n">shuffle_bool</span>


<span class="nd">@check_zip</span>
<span class="k">def</span> <span class="nf">zip</span><span class="p">(</span><span class="n">datasets</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Zip the datasets in the input tuple of datasets.</span>

<span class="sd">    Args:</span>
<span class="sd">        datasets (tuple of class Dataset): A tuple of datasets to be zipped together.</span>
<span class="sd">            The number of datasets must be more than 1.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Dataset, dataset zipped.</span>

<span class="sd">    Raises:</span>
<span class="sd">        ValueError: If the number of datasets is 1.</span>
<span class="sd">        TypeError: If datasets is not a tuple.</span>

<span class="sd">    Examples:</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset which is the combination of dataset_1 and dataset_2</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.zip((dataset_1, dataset_2))</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">datasets</span><span class="p">)</span> <span class="o">&lt;=</span> <span class="mi">1</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;Can&#39;t zip empty or just one dataset!&quot;</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">dataset</span> <span class="ow">in</span> <span class="n">datasets</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">Dataset</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Invalid dataset, expected Dataset object, but got </span><span class="si">%s</span><span class="s2">!&quot;</span> <span class="o">%</span> <span class="nb">type</span><span class="p">(</span><span class="n">dataset</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">ZipDataset</span><span class="p">(</span><span class="n">datasets</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_get_operator_process</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Inner implemented method, mainly for passing sub-process id in C layer</span>

<span class="sd">    Returns:</span>
<span class="sd">         dict, mapping dict of operator id and corresponding process id.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_OP_PROCESS</span>
    <span class="n">process_info</span> <span class="o">=</span> <span class="n">_OP_PROCESS</span>
    <span class="n">op_process</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
    <span class="n">keys</span> <span class="o">=</span> <span class="n">process_info</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span>
    <span class="n">fetched_all</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="n">keys</span><span class="p">:</span>
        <span class="n">op_process</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">process_info</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">1</span><span class="p">])</span>
        <span class="n">item_full</span> <span class="o">=</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">process_info</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">1</span><span class="p">])</span> <span class="o">==</span> <span class="n">process_info</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">0</span><span class="p">])</span>
        <span class="n">fetched_all</span> <span class="o">=</span> <span class="n">fetched_all</span> <span class="ow">and</span> <span class="n">item_full</span>
    <span class="k">return</span> <span class="n">op_process</span><span class="p">,</span> <span class="n">fetched_all</span>


<span class="k">def</span> <span class="nf">_set_dataset_permissions</span><span class="p">(</span><span class="n">file_name</span><span class="p">,</span> <span class="n">num_files</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    set saved dataset files&#39; permissions to 600</span>
<span class="sd">    the rule of dataset filenames should be the same as those in C++.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">num_digits</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">num_files</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))</span>
    <span class="k">if</span> <span class="n">num_files</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">paths</span> <span class="o">=</span> <span class="p">[</span><span class="n">file_name</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">paths</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;</span><span class="si">{}{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">file_name</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">x</span><span class="p">)</span><span class="o">.</span><span class="n">rjust</span><span class="p">(</span><span class="n">num_digits</span><span class="p">,</span> <span class="s1">&#39;0&#39;</span><span class="p">))</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_files</span><span class="p">)]</span>

    <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">paths</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">item</span><span class="p">):</span>
            <span class="n">os</span><span class="o">.</span><span class="n">chmod</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="n">stat</span><span class="o">.</span><span class="n">S_IRUSR</span> <span class="o">|</span> <span class="n">stat</span><span class="o">.</span><span class="n">S_IWUSR</span><span class="p">)</span>
            <span class="n">index_file</span> <span class="o">=</span> <span class="n">item</span> <span class="o">+</span> <span class="s2">&quot;.db&quot;</span>
            <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">index_file</span><span class="p">):</span>
                <span class="n">os</span><span class="o">.</span><span class="n">chmod</span><span class="p">(</span><span class="n">index_file</span><span class="p">,</span> <span class="n">stat</span><span class="o">.</span><span class="n">S_IRUSR</span> <span class="o">|</span> <span class="n">stat</span><span class="o">.</span><span class="n">S_IWUSR</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">Dataset</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a dataset in DataEngine&#39;s data pipeline.</span>

<span class="sd">    This class is the base class of SourceDataset and Dataset, and represents</span>
<span class="sd">    a node in the data flow graph.</span>
<span class="sd">                                     Dataset</span>
<span class="sd">           -----------------------------------------------------------</span>
<span class="sd">           |                  |                   |                  |</span>
<span class="sd">    VisionBaseDataset    TextBaseDataset    AudioBaseDataset         |</span>
<span class="sd">           -                  -                   -                  |</span>
<span class="sd">           |                  |                   |                  |</span>
<span class="sd">           ----------------------------------------                  |</span>
<span class="sd">                      UnionBaseDataset                               |</span>
<span class="sd">                                                                     |</span>
<span class="sd">                                                               SourceDataset</span>
<span class="sd">                                                                     -</span>
<span class="sd">                                                                     |</span>
<span class="sd">                                                              MappableDataset</span>

<span class="sd">    DatasetOperator: MapDataset(UnionBaseDataset)</span>
<span class="sd">                     BatchDataset(UnionBaseDataset)</span>
<span class="sd">                     BucketBatchByLengthDataset(UnionBaseDataset)</span>
<span class="sd">                     ShuffleDataset(UnionBaseDataset)</span>
<span class="sd">                     FilterDataset(UnionBaseDataset)</span>
<span class="sd">                     RepeatDataset(UnionBaseDataset)</span>
<span class="sd">                     SkipDataset(UnionBaseDataset)</span>
<span class="sd">                     TakeDataset(UnionBaseDataset)</span>
<span class="sd">                     ZipDataset(UnionBaseDataset)</span>
<span class="sd">                     ConcatDataset(UnionBaseDataset)</span>
<span class="sd">                     RenameDataset(UnionBaseDataset)</span>
<span class="sd">                     ProjectDataset(UnionBaseDataset)</span>
<span class="sd">                     SyncWaitDataset(UnionBaseDataset)</span>

<span class="sd">    Impl Dataset - vision:       ImageFolderDataset(MappableDataset, VisionBaseDataset)</span>
<span class="sd">                                 USPSDataset(SourceDataset, VisionBaseDataset)</span>
<span class="sd">    Impl Dataset - text:         TextFileDataset(SourceDataset, TextBaseDataset)</span>
<span class="sd">                                 YahooAnswersDataset(SourceDataset, TextBaseDataset)</span>
<span class="sd">    Impl Dataset - audio:        LJSpeechDataset(MappableDataset, AudioBaseDataset)</span>
<span class="sd">                                 TedliumDataset(MappableDataset, AudioBaseDataset)</span>
<span class="sd">    Impl Dataset - standard:     MindDataset(MappableDataset, UnionBaseDataset)</span>
<span class="sd">                                 TFRecordDataset(SourceDataset, UnionBaseDataset)</span>
<span class="sd">    Impl Dataset - user defined: GeneratorDataset(MappableDataset, UnionBaseDataset)</span>
<span class="sd">                                 NumpySlicesDataset(GeneratorDataset)</span>

<span class="sd">    Args:</span>
<span class="sd">        num_parallel_workers (int, optional): Number of workers to process the dataset in parallel</span>
<span class="sd">            (default=None).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="c1"># Note: children and parent are internal variables, not recommended for external using.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">children</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">children</span><span class="p">,</span> <span class="p">[])</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">children</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">children</span> <span class="o">=</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">]</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">child</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="n">child</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">weakref</span><span class="o">.</span><span class="n">ref</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span> <span class="o">=</span> <span class="n">num_parallel_workers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cache</span> <span class="o">=</span> <span class="n">cache</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_device_iter</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_input_indexs</span> <span class="o">=</span> <span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_types</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span> <span class="o">=</span> <span class="p">[</span><span class="kc">False</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_col_names</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_num_classes</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_repeat_count</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_class_indexing</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_sync</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">create_ir_tree</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Internal method to build an IR tree.</span>

<span class="sd">        Returns:</span>
<span class="sd">            DatasetNode, the root node of the IR tree.</span>
<span class="sd">            Dataset, the root dataset of the IR tree.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">parent</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">parent</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">dataset</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="k">global</span> <span class="n">_OP_NAME</span>
        <span class="n">_OP_NAME</span> <span class="o">=</span> <span class="n">Dataset</span><span class="o">.</span><span class="n">_get_operator_id</span><span class="p">(</span><span class="n">dataset</span><span class="p">)</span>
        <span class="n">ir_tree</span> <span class="o">=</span> <span class="n">dataset</span><span class="o">.</span><span class="n">parse_tree</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="n">parent</span>
        <span class="n">_init_device_info</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">ir_tree</span><span class="p">,</span> <span class="n">dataset</span>

    <span class="k">def</span> <span class="nf">close_pool</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Close multiprocessing pool in dataset. If you are familiar with multiprocessing library, you can regard this</span>
<span class="sd">        as a destructor for a processingPool object.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;process_pool&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">child</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="n">child</span><span class="o">.</span><span class="n">close_pool</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">notify_watchdog</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Close watchdog thread in dataset. Now GeneratorDataset/map/batch will use a thread named watch_dog to monitor</span>
<span class="sd">        multiprocess, for get_dataset_size/output_shapes/output_types/get_col_name/num_classes, we need notify_watchdog</span>
<span class="sd">        to close watch_dog thread manually.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;sample_fn&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">sample_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">sample_fn</span><span class="o">.</span><span class="n">multi_process</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">sample_fn</span><span class="o">.</span><span class="n">_abort_watchdog</span><span class="p">()</span>  <span class="c1"># pylint: disable=W0212</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;watch_dog&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;eot&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_abort_watchdog</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">child</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="n">child</span><span class="o">.</span><span class="n">notify_watchdog</span><span class="p">()</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_get_operator_id</span><span class="p">(</span><span class="n">dataset</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Internal method to iterate the tree and obtain op_id of each operator.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, the root dataset of the tree.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">op_name</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
        <span class="n">generator_process</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
        <span class="n">op_name</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">dataset</span><span class="p">)]</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="n">op_id</span> <span class="o">=</span> <span class="mi">1</span>

        <span class="k">def</span> <span class="nf">process_name</span><span class="p">(</span><span class="n">datasets</span><span class="p">,</span> <span class="n">operator_id</span><span class="p">):</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">datasets</span><span class="p">:</span>
                <span class="k">return</span> <span class="mi">0</span>
            <span class="n">temp</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">datasets</span><span class="p">:</span>
                <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="n">item</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
                    <span class="n">temp</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">d</span><span class="p">)</span>
                    <span class="n">op_name</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">d</span><span class="p">)]</span> <span class="o">=</span> <span class="n">operator_id</span>

                    <span class="kn">from</span> <span class="nn">mindspore.dataset.engine.datasets_user_defined</span> <span class="kn">import</span> <span class="n">GeneratorDataset</span>
                    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">d</span><span class="p">,</span> <span class="n">GeneratorDataset</span><span class="p">)</span> <span class="ow">and</span> <span class="n">d</span><span class="o">.</span><span class="n">sample_fn</span> <span class="ow">and</span> <span class="n">d</span><span class="o">.</span><span class="n">sample_fn</span><span class="o">.</span><span class="n">pids</span><span class="p">:</span>
                        <span class="n">generator_process</span><span class="p">[</span><span class="n">operator_id</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="n">d</span><span class="o">.</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="nb">set</span><span class="p">(</span><span class="n">d</span><span class="o">.</span><span class="n">sample_fn</span><span class="o">.</span><span class="n">pids</span><span class="p">)]</span>

            <span class="n">operator_id</span> <span class="o">=</span> <span class="n">operator_id</span> <span class="o">+</span> <span class="mi">1</span>
            <span class="k">return</span> <span class="n">process_name</span><span class="p">(</span><span class="n">temp</span><span class="p">,</span> <span class="n">operator_id</span><span class="p">)</span>

        <span class="n">process_name</span><span class="p">([</span><span class="n">dataset</span><span class="p">],</span> <span class="n">op_id</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">generator_process</span><span class="p">:</span>
            <span class="k">global</span> <span class="n">_OP_PROCESS</span>
            <span class="n">_OP_PROCESS</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">generator_process</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">op_name</span>

    <span class="k">def</span> <span class="nf">parse_tree</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Internal method to parse the API tree into an IR tree.</span>

<span class="sd">        Returns:</span>
<span class="sd">            DatasetNode, the root node of the IR tree.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">parent</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The data pipeline is not a tree (i.e., one node has 2 consumers)&quot;</span><span class="p">)</span>
        <span class="n">ir_children</span> <span class="o">=</span> <span class="p">[</span><span class="n">d</span><span class="o">.</span><span class="n">parse_tree</span><span class="p">()</span> <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">]</span>
        <span class="c1"># Bootstrap can only be performed on a copy of the original dataset node.</span>
        <span class="c1"># Bootstrap on original dataset node will make all iterators share the same process pool</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">iterator_bootstrap</span><span class="p">()</span>
        <span class="n">ir_node</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">parse</span><span class="p">(</span><span class="n">ir_children</span><span class="p">)</span>
        <span class="n">ir_node</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">post_parse</span><span class="p">(</span><span class="n">ir_node</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">ir_node</span>

    <span class="k">def</span> <span class="nf">__safe_deepcopy__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">memodict</span><span class="p">,</span> <span class="n">exclude</span><span class="o">=</span><span class="p">()):</span>
        <span class="k">if</span> <span class="nb">id</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="ow">in</span> <span class="n">memodict</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">memodict</span><span class="p">[</span><span class="nb">id</span><span class="p">(</span><span class="bp">self</span><span class="p">)]</span>
        <span class="bp">cls</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span>
        <span class="n">new_op</span> <span class="o">=</span> <span class="bp">cls</span><span class="o">.</span><span class="fm">__new__</span><span class="p">(</span><span class="bp">cls</span><span class="p">)</span>
        <span class="n">memodict</span><span class="p">[</span><span class="nb">id</span><span class="p">(</span><span class="bp">self</span><span class="p">)]</span> <span class="o">=</span> <span class="n">new_op</span>
        <span class="k">for</span> <span class="n">arg</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="k">if</span> <span class="n">arg</span> <span class="ow">in</span> <span class="n">exclude</span><span class="p">:</span>
                <span class="nb">setattr</span><span class="p">(</span><span class="n">new_op</span><span class="p">,</span> <span class="n">arg</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">try</span><span class="p">:</span>
                    <span class="nb">setattr</span><span class="p">(</span><span class="n">new_op</span><span class="p">,</span> <span class="n">arg</span><span class="p">,</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">memodict</span><span class="p">))</span>
                <span class="k">except</span> <span class="ne">TypeError</span><span class="p">:</span>
                    <span class="nb">setattr</span><span class="p">(</span><span class="n">new_op</span><span class="p">,</span> <span class="n">arg</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">new_op</span>

    <span class="k">def</span> <span class="nf">iterator_bootstrap</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">pass</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_noop_mode</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">_is_role_sched</span><span class="p">()</span> <span class="ow">or</span> <span class="n">_is_role_pserver</span><span class="p">():</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="fm">__add__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">datasets</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">concat</span><span class="p">(</span><span class="n">datasets</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">to_json</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Serialize a pipeline into JSON string and dump into file if filename is provided.</span>

<span class="sd">        Args:</span>
<span class="sd">            filename (str): filename of JSON file to be saved as (default=&quot;&quot;).</span>

<span class="sd">        Returns:</span>
<span class="sd">            str, JSON string of the pipeline.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ir_tree</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">json</span><span class="o">.</span><span class="n">loads</span><span class="p">(</span><span class="n">ir_tree</span><span class="o">.</span><span class="n">to_json</span><span class="p">(</span><span class="n">filename</span><span class="p">))</span>

    <span class="nd">@check_bucket_batch_by_length</span>
    <span class="k">def</span> <span class="nf">bucket_batch_by_length</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">column_names</span><span class="p">,</span> <span class="n">bucket_boundaries</span><span class="p">,</span> <span class="n">bucket_batch_sizes</span><span class="p">,</span> <span class="n">element_length_function</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                               <span class="n">pad_info</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">pad_to_bucket_boundary</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">drop_remainder</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Bucket elements according to their lengths. Each bucket will be padded and batched when</span>
<span class="sd">        they are full.</span>

<span class="sd">        A length function is called on each row in the dataset. The row is then</span>
<span class="sd">        bucketed based on its length and bucket boundaries. When a bucket reaches its</span>
<span class="sd">        corresponding size specified in bucket_batch_sizes, the entire bucket will be</span>
<span class="sd">        padded according to pad_info, and then form a batch.</span>
<span class="sd">        Each batch will be full, except one special case: the last batch for each bucket may not be full.</span>

<span class="sd">        Args:</span>
<span class="sd">            column_names (list[str]): Columns passed to element_length_function.</span>
<span class="sd">            bucket_boundaries (list[int]): A list consisting of the upper boundaries</span>
<span class="sd">                of the buckets. Must be strictly increasing. If there are n boundaries,</span>
<span class="sd">                n+1 buckets are created: One bucket for [0, bucket_boundaries[0]), one</span>
<span class="sd">                bucket for [bucket_boundaries[i], bucket_boundaries[i+1]) for each</span>
<span class="sd">                0&lt;i&lt;n-1, and the last bucket for [bucket_boundaries[n-1], inf).</span>
<span class="sd">            bucket_batch_sizes (list[int]): A list consisting of the batch sizes for</span>
<span class="sd">                each bucket. Must contain len(bucket_boundaries)+1 elements.</span>
<span class="sd">            element_length_function (Callable, optional): A function that takes in</span>
<span class="sd">                M arguments where M = len(column_names) and returns an integer. If no value</span>
<span class="sd">                provided, parameter M the len(column_names) must be 1, and the size of the first</span>
<span class="sd">                dimension of that column will be taken as the length (default=None).</span>
<span class="sd">            pad_info (dict, optional): The information about how to batch each column. The key</span>
<span class="sd">                corresponds to the column name, and the value must be a tuple of 2 elements.</span>
<span class="sd">                The first element corresponds to the shape to pad to, and the second</span>
<span class="sd">                element corresponds to the value to pad with. If a column is not</span>
<span class="sd">                specified, then that column will be padded to the longest in the current</span>
<span class="sd">                batch, and 0 will be used as the padding value. Any None dimensions will</span>
<span class="sd">                be padded to the longest in the current batch, unless if</span>
<span class="sd">                pad_to_bucket_boundary is True. If no padding is wanted, set pad_info</span>
<span class="sd">                to None (default=None).</span>
<span class="sd">            pad_to_bucket_boundary (bool, optional): If True, will pad each None</span>
<span class="sd">                dimension in pad_info to the bucket_boundary minus 1. If there are any</span>
<span class="sd">                elements that fall into the last bucket, an error will occur</span>
<span class="sd">                (default=False).</span>
<span class="sd">            drop_remainder (bool, optional): If True, will drop the last batch for each</span>
<span class="sd">                bucket if it is not a full batch (default=False).</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset bucketed and batched by length.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where certain counts rows are combined into a batch</span>
<span class="sd">            &gt;&gt;&gt; # and drops the last incomplete batch if there is one.</span>
<span class="sd">            &gt;&gt;&gt; import numpy as np</span>
<span class="sd">            &gt;&gt;&gt; def generate_2_columns(n):</span>
<span class="sd">            ...     for i in range(n):</span>
<span class="sd">            ...         yield (np.array([i]), np.array([j for j in range(i + 1)]))</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; column_names = [&quot;col1&quot;, &quot;col2&quot;]</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.GeneratorDataset(generate_2_columns(8), column_names)</span>
<span class="sd">            &gt;&gt;&gt; bucket_boundaries = [5, 10]</span>
<span class="sd">            &gt;&gt;&gt; bucket_batch_sizes = [2, 1, 1]</span>
<span class="sd">            &gt;&gt;&gt; element_length_function = (lambda col1, col2: max(len(col1), len(col2)))</span>
<span class="sd">            &gt;&gt;&gt; # Will pad col2 to shape [bucket_boundaries[i]] where i is the</span>
<span class="sd">            &gt;&gt;&gt; # index of the bucket that is currently being batched.</span>
<span class="sd">            &gt;&gt;&gt; pad_info = {&quot;col2&quot;: ([None], -1)}</span>
<span class="sd">            &gt;&gt;&gt; pad_to_bucket_boundary = True</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.bucket_batch_by_length(column_names, bucket_boundaries,</span>
<span class="sd">            ...                                          bucket_batch_sizes,</span>
<span class="sd">            ...                                          element_length_function, pad_info,</span>
<span class="sd">            ...                                          pad_to_bucket_boundary)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">BucketBatchByLengthDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">column_names</span><span class="p">,</span> <span class="n">bucket_boundaries</span><span class="p">,</span> <span class="n">bucket_batch_sizes</span><span class="p">,</span>
                                          <span class="n">element_length_function</span><span class="p">,</span> <span class="n">pad_info</span><span class="p">,</span> <span class="n">pad_to_bucket_boundary</span><span class="p">,</span> <span class="n">drop_remainder</span><span class="p">)</span>

    <span class="nd">@check_batch</span>
    <span class="k">def</span> <span class="nf">batch</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">drop_remainder</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">per_batch_map</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
              <span class="n">input_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">output_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">column_order</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">pad_info</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
              <span class="n">python_multiprocessing</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="mi">16</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Combine batch_size number of consecutive rows into batches.</span>

<span class="sd">        For any child node, a batch is treated as a single row.</span>
<span class="sd">        For any column, all the elements within that column must have the same shape.</span>
<span class="sd">        If a per_batch_map callable is provided, it will be applied to the batches of tensors.</span>

<span class="sd">        Note:</span>
<span class="sd">            The order of using repeat and batch reflects the number of batches and per_batch_map.</span>
<span class="sd">            It is recommended that the repeat operation applied after the batch operation finished.</span>

<span class="sd">        Args:</span>
<span class="sd">            batch_size (int or function): The number of rows each batch is created with. An</span>
<span class="sd">                int or callable object which takes exactly 1 parameter, BatchInfo.</span>
<span class="sd">            drop_remainder (bool, optional): Determines whether or not to drop the last block</span>
<span class="sd">                whose data row number is less than batch size (default=False). If True, and if there are less</span>
<span class="sd">                than batch_size rows available to make the last batch, then those rows will</span>
<span class="sd">                be dropped and not propagated to the child node.</span>
<span class="sd">            num_parallel_workers (int, optional): Number of workers(threads) to process the dataset in parallel</span>
<span class="sd">                (default=None).</span>
<span class="sd">            per_batch_map (callable, optional): Per batch map callable (default=None). A callable which takes</span>
<span class="sd">                (list[Tensor], list[Tensor], ..., BatchInfo) as input parameters. Each list[Tensor] represents a batch</span>
<span class="sd">                of Tensors on a given column. The number of lists should match with the number of entries in</span>
<span class="sd">                input_columns. The last parameter of the callable should always be a BatchInfo object. Per_batch_map</span>
<span class="sd">                should return (list[Tensor], list[Tensor], ...). The length of each list in output should be the same as</span>
<span class="sd">                the input. output_columns is required if the number of output lists is different from input.</span>
<span class="sd">            input_columns (Union[str, list[str]], optional): List of names of the input columns. The size of the list</span>
<span class="sd">                should match with signature of per_batch_map callable (default=None).</span>
<span class="sd">            output_columns (Union[str, list[str]], optional): List of names assigned to the columns</span>
<span class="sd">                outputted by the last operation. This parameter is mandatory if len(input_columns) !=</span>
<span class="sd">                len(output_columns). The size of this list must match the number of output</span>
<span class="sd">                columns of the last operation. (default=None, output columns will have the same</span>
<span class="sd">                name as the input columns, i.e., the columns will be replaced).</span>
<span class="sd">            column_order (Union[str, list[str]], optional): Specifies the list of all the columns you need in the whole</span>
<span class="sd">                dataset (default=None). The parameter is required when len(input_column) != len(output_column).</span>
<span class="sd">                Caution: the list here is not just the columns specified in parameter input_columns and output_columns.</span>
<span class="sd">            pad_info (dict, optional): Whether to perform padding on selected columns. pad_info={&quot;col1&quot;:([224,224],0)}</span>
<span class="sd">                would pad column with name &quot;col1&quot; to a tensor of size [224,224] and fill the missing with 0</span>
<span class="sd">                (default=None).</span>
<span class="sd">            python_multiprocessing (bool, optional): Parallelize Python function per_batch_map with multi-processing.</span>
<span class="sd">                This option could be beneficial if the function is computational heavy (default=False).</span>
<span class="sd">            max_rowsize(int, optional): Maximum size of row in MB that is used for shared memory allocation to copy</span>
<span class="sd">               data between processes.  This is only used if python_multiprocessing is set to True (default=16).</span>

<span class="sd">        Returns:</span>
<span class="sd">            BatchDataset, dataset batched.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where every 100 rows are combined into a batch</span>
<span class="sd">            &gt;&gt;&gt; # and drops the last incomplete batch if there is one.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.batch(100, True)</span>
<span class="sd">            &gt;&gt;&gt; # resize image according to its batch number, if it&#39;s 5-th batch, resize to (5^2, 5^2) = (25, 25)</span>
<span class="sd">            &gt;&gt;&gt; def np_resize(col, batchInfo):</span>
<span class="sd">            ...     output = col.copy()</span>
<span class="sd">            ...     s = (batchInfo.get_batch_num() + 1) ** 2</span>
<span class="sd">            ...     index = 0</span>
<span class="sd">            ...     for c in col:</span>
<span class="sd">            ...         img = Image.fromarray(c.astype(&#39;uint8&#39;)).convert(&#39;RGB&#39;)</span>
<span class="sd">            ...         img = img.resize((s, s), Image.ANTIALIAS)</span>
<span class="sd">            ...         output[index] = np.array(img)</span>
<span class="sd">            ...         index += 1</span>
<span class="sd">            ...     return (output,)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.batch(batch_size=8, input_columns=[&quot;image&quot;], per_batch_map=np_resize)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">BatchDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">drop_remainder</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">per_batch_map</span><span class="p">,</span> <span class="n">input_columns</span><span class="p">,</span>
                            <span class="n">output_columns</span><span class="p">,</span> <span class="n">column_order</span><span class="p">,</span> <span class="n">pad_info</span><span class="p">,</span> <span class="n">python_multiprocessing</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="p">)</span>

    <span class="nd">@check_sync_wait</span>
    <span class="k">def</span> <span class="nf">sync_wait</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">condition_name</span><span class="p">,</span> <span class="n">num_batch</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">callback</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Add a blocking condition to the input Dataset. A synchronize action will be applied.</span>

<span class="sd">        Args:</span>
<span class="sd">            condition_name (str): The condition name that is used to toggle sending next row.</span>
<span class="sd">            num_batch (int): the number of batches without blocking at the start of each epoch (default=1).</span>
<span class="sd">            callback (function): The callback function that will be invoked when sync_update is called (default=None).</span>

<span class="sd">        Returns:</span>
<span class="sd">            SyncWaitDataset, dataset added a blocking condition.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: If condition name already exists.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; import numpy as np</span>
<span class="sd">            &gt;&gt;&gt; def gen():</span>
<span class="sd">            ...     for i in range(100):</span>
<span class="sd">            ...         yield (np.array(i),)</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; class Augment:</span>
<span class="sd">            ...     def __init__(self, loss):</span>
<span class="sd">            ...         self.loss = loss</span>
<span class="sd">            ...</span>
<span class="sd">            ...     def preprocess(self, input_):</span>
<span class="sd">            ...         return input_</span>
<span class="sd">            ...</span>
<span class="sd">            ...     def update(self, data):</span>
<span class="sd">            ...         self.loss = data[&quot;loss&quot;]</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; batch_size = 4</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.GeneratorDataset(gen, column_names=[&quot;input&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; aug = Augment(0)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.sync_wait(condition_name=&quot;policy&quot;, callback=aug.update)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[aug.preprocess], input_columns=[&quot;input&quot;])</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.batch(batch_size)</span>
<span class="sd">            &gt;&gt;&gt; count = 0</span>
<span class="sd">            &gt;&gt;&gt; for data in dataset.create_dict_iterator(num_epochs=1, output_numpy=True):</span>
<span class="sd">            ...     assert data[&quot;input&quot;][0] == count</span>
<span class="sd">            ...     count += batch_size</span>
<span class="sd">            ...     data = {&quot;loss&quot;: count}</span>
<span class="sd">            ...     dataset.sync_update(condition_name=&quot;policy&quot;, data=data)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">SyncWaitDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">condition_name</span><span class="p">,</span> <span class="n">num_batch</span><span class="p">,</span> <span class="n">callback</span><span class="p">)</span>

    <span class="nd">@check_shuffle</span>
    <span class="k">def</span> <span class="nf">shuffle</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Randomly shuffles the rows of this dataset using the following policy:</span>

<span class="sd">        1. Make a shuffle buffer that contains the first buffer_size rows.</span>
<span class="sd">        2. Randomly select an element from the shuffle buffer to be the next row</span>
<span class="sd">           propagated to the child node.</span>
<span class="sd">        3. Get the next row (if any) from the parent node and put it in the shuffle buffer.</span>
<span class="sd">        4. Repeat steps 2 and 3 until there are no more rows left in the shuffle buffer.</span>

<span class="sd">        A random seed can be provided to be used on the first epoch. In every subsequent</span>
<span class="sd">        epoch, the seed is changed to a new one, randomly generated value.</span>

<span class="sd">        Args:</span>
<span class="sd">            buffer_size (int): The size of the buffer (must be larger than 1) for</span>
<span class="sd">                shuffling. Setting buffer_size equal to the number of rows in the entire</span>
<span class="sd">                dataset will result in a global shuffle.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset shuffled.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: If exist sync operators before shuffle.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; # Optionally set the seed for the first epoch</span>
<span class="sd">            &gt;&gt;&gt; ds.config.set_seed(58)</span>
<span class="sd">            &gt;&gt;&gt; # Create a shuffled dataset using a shuffle buffer of size 4</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.shuffle(4)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">ShuffleDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">flat_map</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">func</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Map `func` to each row in dataset and flatten the result.</span>

<span class="sd">        The specified `func` is a function that must take one &#39;Ndarray&#39; as input</span>
<span class="sd">        and return a &#39;Dataset&#39;.</span>

<span class="sd">        Args:</span>
<span class="sd">            func (function): A function that must take one &#39;Ndarray&#39; as an argument and</span>
<span class="sd">                return a &#39;Dataset&#39;.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset applied by the function.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # use NumpySlicesDataset as an example</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.NumpySlicesDataset([[0, 1], [2, 3]])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; def flat_map_func(array):</span>
<span class="sd">            ...     # create a NumpySlicesDataset with the array</span>
<span class="sd">            ...     dataset = ds.NumpySlicesDataset(array)</span>
<span class="sd">            ...     # repeat the dataset twice</span>
<span class="sd">            ...     dataset = dataset.repeat(2)</span>
<span class="sd">            ...     return dataset</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.flat_map(flat_map_func)</span>
<span class="sd">            &gt;&gt;&gt; # [[0, 1], [0, 1], [2, 3], [2, 3]]</span>

<span class="sd">        Raises:</span>
<span class="sd">            TypeError: If `func` is not a function.</span>
<span class="sd">            TypeError: If `func` doesn&#39;t return a Dataset.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">dataset</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="s1">&#39;__call__&#39;</span><span class="p">):</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">critical</span><span class="p">(</span><span class="s2">&quot;func must be a function.&quot;</span><span class="p">)</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;func must be a function.&quot;</span><span class="p">)</span>

        <span class="k">for</span> <span class="n">row_data</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_tuple_iterator</span><span class="p">(</span><span class="n">output_numpy</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">dataset</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dataset</span> <span class="o">=</span> <span class="n">func</span><span class="p">(</span><span class="n">row_data</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dataset</span> <span class="o">+=</span> <span class="n">func</span><span class="p">(</span><span class="n">row_data</span><span class="p">)</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">Dataset</span><span class="p">):</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">critical</span><span class="p">(</span><span class="s2">&quot;flat_map must return a Dataset object.&quot;</span><span class="p">)</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;flat_map must return a Dataset object.&quot;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">dataset</span>

    <span class="nd">@check_map</span>
    <span class="k">def</span> <span class="nf">map</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">operations</span><span class="p">,</span> <span class="n">input_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">output_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">column_order</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
            <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">python_multiprocessing</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">callbacks</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
            <span class="n">max_rowsize</span><span class="o">=</span><span class="mi">16</span><span class="p">,</span> <span class="n">offload</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Apply each operation in operations to this dataset.</span>

<span class="sd">        The order of operations is determined by the position of each operation in the operations parameter.</span>
<span class="sd">        operations[0] will be applied first, then operations[1], then operations[2], etc.</span>

<span class="sd">        Each operation will be passed one or more columns from the dataset as input, and zero or</span>
<span class="sd">        more columns will be outputted. The first operation will be passed the columns specified</span>
<span class="sd">        in input_columns as input. If there is more than one operator in operations, the outputted</span>
<span class="sd">        columns of the previous operation are used as the input columns for the next operation.</span>
<span class="sd">        The columns outputted by the very last operation will be assigned names specified by</span>
<span class="sd">        output_columns.</span>

<span class="sd">        Only the columns specified in column_order will be propagated to the child node. These</span>
<span class="sd">        columns will be in the same order as specified in column_order.</span>

<span class="sd">        Args:</span>
<span class="sd">            operations (Union[list[TensorOp], list[functions]]): List of operations to be</span>
<span class="sd">                applied on the dataset. Operations are applied in the order they appear in this list.</span>
<span class="sd">            input_columns (Union[str, list[str]], optional): List of the names of the columns that will be passed to</span>
<span class="sd">                the first operation as input. The size of this list must match the number of</span>
<span class="sd">                input columns expected by the first operator. (default=None, the first</span>
<span class="sd">                operation will be passed however many columns that are required, starting from</span>
<span class="sd">                the first column).</span>
<span class="sd">            output_columns (Union[str, list[str]], optional): List of names assigned to the columns outputted by</span>
<span class="sd">                the last operation. This parameter is mandatory if len(input_columns) !=</span>
<span class="sd">                len(output_columns). The size of this list must match the number of output</span>
<span class="sd">                columns of the last operation. (default=None, output columns will have the same</span>
<span class="sd">                name as the input columns, i.e., the columns will be replaced).</span>
<span class="sd">            column_order (list[str], optional): Specifies the list of all the columns you need in the whole</span>
<span class="sd">                dataset (default=None). The parameter is required when len(input_column) != len(output_column).</span>
<span class="sd">                Caution: the list here is not just the columns specified in parameter input_columns and output_columns.</span>
<span class="sd">            num_parallel_workers (int, optional): Number of threads used to process the dataset in</span>
<span class="sd">                parallel (default=None, the value from the configuration will be used).</span>
<span class="sd">            python_multiprocessing (bool, optional): Parallelize Python operations with multiple worker processes. This</span>
<span class="sd">                option could be beneficial if the Python operation is computational heavy (default=False).</span>
<span class="sd">            cache (DatasetCache, optional): Use tensor caching service to speed up dataset processing.</span>
<span class="sd">                (default=None, which means no cache is used).</span>
<span class="sd">            callbacks (DSCallback, list[DSCallback], optional): List of Dataset callbacks to be called (Default=None).</span>
<span class="sd">            max_rowsize (int, optional): Maximum size of row in MB that is used for shared memory allocation to copy</span>
<span class="sd">               data between processes.  This is only used if python_multiprocessing is set to True (Default=16).</span>
<span class="sd">            offload (bool, optional): Flag to indicate whether offload is used (Default=None).</span>

<span class="sd">        Note:</span>
<span class="sd">            - Input `operations` mainly accept c_transforms, py_transforms operator in mindspore.dataset part, plus user</span>
<span class="sd">              defined Python function(PyFuncs).</span>
<span class="sd">            - Do not add network computing operators from mindspore.nn and mindspore.ops or others into this</span>
<span class="sd">              `operations`.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset after mapping operation.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance of Dataset which has 2 columns, &quot;image&quot; and &quot;label&quot;.</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Define two operations, where each operation accepts 1 input column and outputs 1 column.</span>
<span class="sd">            &gt;&gt;&gt; decode_op = c_vision.Decode(rgb=True)</span>
<span class="sd">            &gt;&gt;&gt; random_jitter_op = c_vision.RandomColorAdjust(brightness=(0.8, 0.8), contrast=(1, 1),</span>
<span class="sd">            ...                                               saturation=(1, 1), hue=(0, 0))</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # 1) Simple map example.</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Apply decode_op on column &quot;image&quot;. This column will be replaced by the outputted</span>
<span class="sd">            &gt;&gt;&gt; # column of decode_op. Since column_order is not provided, both columns &quot;image&quot;</span>
<span class="sd">            &gt;&gt;&gt; # and &quot;label&quot; will be propagated to the child node in their original order.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op], input_columns=[&quot;image&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Decode and rename column &quot;image&quot; to &quot;decoded_image&quot;.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op], input_columns=[&quot;image&quot;], output_columns=[&quot;decoded_image&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Specify the order of the output columns.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op], input_columns=[&quot;image&quot;],</span>
<span class="sd">            ...                       output_columns=None, column_order=[&quot;label&quot;, &quot;image&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Rename column &quot;image&quot; to &quot;decoded_image&quot; and also specify the order of the output columns.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op], input_columns=[&quot;image&quot;],</span>
<span class="sd">            ...                       output_columns=[&quot;decoded_image&quot;], column_order=[&quot;label&quot;, &quot;decoded_image&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Rename column &quot;image&quot; to &quot;decoded_image&quot; and keep only this column.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op], input_columns=[&quot;image&quot;],</span>
<span class="sd">            ...                       output_columns=[&quot;decoded_image&quot;], column_order=[&quot;decoded_image&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # A simple example for mapping pyfunc. Renaming columns and specifying column order</span>
<span class="sd">            &gt;&gt;&gt; # work in the same way as the previous examples.</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.NumpySlicesDataset(data=[[0, 1, 2]], column_names=[&quot;data&quot;])</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[(lambda x: x + 1)], input_columns=[&quot;data&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # 2) Map example with more than one operation.</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where the images are decoded, then randomly color jittered.</span>
<span class="sd">            &gt;&gt;&gt; # decode_op takes column &quot;image&quot; as input and outputs one column. The column</span>
<span class="sd">            &gt;&gt;&gt; # outputted by decode_op is passed as input to random_jitter_op.</span>
<span class="sd">            &gt;&gt;&gt; # random_jitter_op will output one column. Column &quot;image&quot; will be replaced by</span>
<span class="sd">            &gt;&gt;&gt; # the column outputted by random_jitter_op (the very last operation). All other</span>
<span class="sd">            &gt;&gt;&gt; # columns are unchanged. Since column_order is not specified, the order of the</span>
<span class="sd">            &gt;&gt;&gt; # columns will remain the same.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op, random_jitter_op], input_columns=[&quot;image&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Rename the column outputted by random_jitter_op to &quot;image_mapped&quot;.</span>
<span class="sd">            &gt;&gt;&gt; # Specifying column order works in the same way as examples in 1).</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[decode_op, random_jitter_op], input_columns=[&quot;image&quot;],</span>
<span class="sd">            ...                       output_columns=[&quot;image_mapped&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Map with multiple operations using pyfunc. Renaming columns and specifying column order</span>
<span class="sd">            &gt;&gt;&gt; # work in the same way as examples in 1).</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.NumpySlicesDataset(data=[[0, 1, 2]], column_names=[&quot;data&quot;])</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations=[(lambda x: x * x), (lambda x: x - 1)], input_columns=[&quot;data&quot;],</span>
<span class="sd">            ...                                   output_columns=[&quot;data_mapped&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # 3) Example where number of input columns is not equal to number of output columns.</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # operations[0] is a lambda that takes 2 columns as input and outputs 3 columns.</span>
<span class="sd">            &gt;&gt;&gt; # operations[1] is a lambda that takes 3 columns as input and outputs 1 column.</span>
<span class="sd">            &gt;&gt;&gt; # operations[2] is a lambda that takes 1 column as input and outputs 4 columns.</span>
<span class="sd">            &gt;&gt;&gt; #</span>
<span class="sd">            &gt;&gt;&gt; # Note: The number of output columns of operation[i] must equal the number of</span>
<span class="sd">            &gt;&gt;&gt; # input columns of operation[i+1]. Otherwise, this map call will also result</span>
<span class="sd">            &gt;&gt;&gt; # in an error.</span>
<span class="sd">            &gt;&gt;&gt; operations = [(lambda x, y: (x, x + y, x + y + 1)),</span>
<span class="sd">            ...               (lambda x, y, z: x * y * z),</span>
<span class="sd">            ...               (lambda x: (x % 2, x % 3, x % 5, x % 7))]</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Note: Since the number of input columns is not the same as the number of</span>
<span class="sd">            &gt;&gt;&gt; # output columns, the output_columns and column_order parameters must be</span>
<span class="sd">            &gt;&gt;&gt; # specified. Otherwise, this map call will also result in an error.</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.NumpySlicesDataset(data=([[0, 1, 2]], [[3, 4, 5]]), column_names=[&quot;x&quot;, &quot;y&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Propagate all columns to the child node in this order:</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations, input_columns=[&quot;x&quot;, &quot;y&quot;],</span>
<span class="sd">            ...                       output_columns=[&quot;mod2&quot;, &quot;mod3&quot;, &quot;mod5&quot;, &quot;mod7&quot;],</span>
<span class="sd">            ...                       column_order=[&quot;mod2&quot;, &quot;mod3&quot;, &quot;mod5&quot;, &quot;mod7&quot;])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Propagate some columns to the child node in this order:</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.map(operations, input_columns=[&quot;x&quot;, &quot;y&quot;],</span>
<span class="sd">            ...                       output_columns=[&quot;mod2&quot;, &quot;mod3&quot;, &quot;mod5&quot;, &quot;mod7&quot;],</span>
<span class="sd">            ...                       column_order=[&quot;mod7&quot;, &quot;mod3&quot;, &quot;col2&quot;])</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;operator_mixed&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;operator_mixed&#39;</span><span class="p">)</span> <span class="ow">is</span> <span class="kc">True</span><span class="p">:</span>
            <span class="n">num_parallel_workers</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                <span class="s2">&quot;Input &#39;operations&#39; of &#39;map&#39; includes network computing operators like in mindspore.nn, mindspore.ops, &quot;</span>
                <span class="s2">&quot;mindspore.numpy module and etc, which do not support multi-thread compiling, recommend to replace it &quot;</span>
                <span class="s2">&quot;with python implemented operator like numpy etc. Here decrease &#39;num_parallel_workers&#39; into 1.&quot;</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">MapDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">operations</span><span class="p">,</span> <span class="n">input_columns</span><span class="p">,</span> <span class="n">output_columns</span><span class="p">,</span> <span class="n">column_order</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="p">,</span>
                          <span class="n">python_multiprocessing</span><span class="p">,</span> <span class="n">cache</span><span class="p">,</span> <span class="n">callbacks</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="p">,</span> <span class="n">offload</span><span class="p">)</span>

    <span class="nd">@check_filter</span>
    <span class="k">def</span> <span class="nf">filter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predicate</span><span class="p">,</span> <span class="n">input_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Filter dataset by prediction.</span>

<span class="sd">        Note:</span>
<span class="sd">             If input_columns not provided or provided with empty, all columns will be used.</span>

<span class="sd">        Args:</span>
<span class="sd">            predicate (callable): Python callable which returns a boolean value. If False then filter the element.</span>
<span class="sd">            input_columns (Union[str, list[str]], optional): List of names of the input columns, when</span>
<span class="sd">                default=None, the predicate will be applied on all columns in the dataset.</span>
<span class="sd">            num_parallel_workers (int, optional): Number of workers to process the dataset</span>
<span class="sd">                in parallel (default=None).</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset filtered.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # generator data(0 ~ 63)</span>
<span class="sd">            &gt;&gt;&gt; # filter the data that greater than or equal to 11</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.filter(predicate=lambda data: data &lt; 11, input_columns = [&quot;data&quot;])</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">FilterDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predicate</span><span class="p">,</span> <span class="n">input_columns</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="p">)</span>

    <span class="nd">@check_repeat</span>
    <span class="k">def</span> <span class="nf">repeat</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Repeat this dataset `count` times. Repeat infinitely if the count is None or -1.</span>

<span class="sd">        Note:</span>
<span class="sd">            The order of using repeat and batch reflects the number of batches. It is recommended that</span>
<span class="sd">            the repeat operation is used after the batch operation.</span>

<span class="sd">        Args:</span>
<span class="sd">            count (int): Number of times the dataset is going to be repeated (default=None).</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset repeated.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where the dataset is repeated for 50 epochs</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.repeat(50)</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where each epoch is shuffled individually</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.shuffle(10)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.repeat(50)</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where the dataset is first repeated for</span>
<span class="sd">            &gt;&gt;&gt; # 50 epochs before shuffling. The shuffle operator will treat</span>
<span class="sd">            &gt;&gt;&gt; # the entire 50 epochs as one big dataset.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.repeat(50)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.shuffle(10)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">RepeatDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">)</span>

    <span class="nd">@check_skip</span>
    <span class="k">def</span> <span class="nf">skip</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Skip the first N elements of this dataset.</span>

<span class="sd">        Args:</span>
<span class="sd">            count (int): Number of elements in the dataset to be skipped.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset that containing rows like origin rows subtract skipped rows.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset which skips first 3 elements from data</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.skip(3)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">SkipDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">)</span>

    <span class="nd">@check_take</span>
    <span class="k">def</span> <span class="nf">take</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="o">=-</span><span class="mi">1</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Takes at most given numbers of elements from the dataset.</span>

<span class="sd">        Note:</span>
<span class="sd">            1. If count is greater than the number of elements in the dataset or equal to -1,</span>
<span class="sd">               all the elements in dataset will be taken.</span>
<span class="sd">            2. The order of using take and batch matters. If take is before batch operation,</span>
<span class="sd">               then take the given number of rows; otherwise take the given number of batches.</span>

<span class="sd">        Args:</span>
<span class="sd">            count (int, optional): Number of elements to be taken from the dataset (default=-1).</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset taken.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where the dataset includes 50 elements.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.take(50)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">TakeDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_get_absolute_split_sizes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sizes</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Internal method called by split to calculate absolute split sizes and to</span>
<span class="sd">        do some error checking after calculating absolute split sizes.</span>

<span class="sd">        Returns:</span>
<span class="sd">            int, absolute split sizes of the dataset.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># Call get_dataset_size here and check input here because</span>
        <span class="c1"># don&#39;t want to call this once in check_split and another time in</span>
        <span class="c1"># here again</span>
        <span class="n">dataset_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_dataset_size</span><span class="p">()</span>

        <span class="k">if</span> <span class="n">dataset_size</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">dataset_size</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;dataset_size is unknown, unable to split.&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">sizes</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;sizes must be a list.&quot;</span><span class="p">)</span>

        <span class="n">all_int</span> <span class="o">=</span> <span class="nb">all</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">sizes</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">all_int</span><span class="p">:</span>
            <span class="n">sizes_sum</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">sizes_sum</span> <span class="o">!=</span> <span class="n">dataset_size</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Sum of split sizes </span><span class="si">{}</span><span class="s2"> is not equal to dataset size </span><span class="si">{}</span><span class="s2">.&quot;</span>
                                   <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">sizes_sum</span><span class="p">,</span> <span class="n">dataset_size</span><span class="p">))</span>
            <span class="k">return</span> <span class="n">sizes</span>

        <span class="n">absolute_sizes</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">sizes</span><span class="p">:</span>
            <span class="n">absolute_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">item</span> <span class="o">*</span> <span class="n">dataset_size</span><span class="p">))</span>
            <span class="k">if</span> <span class="n">absolute_size</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Split percentage </span><span class="si">{}</span><span class="s2"> is too small.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">item</span><span class="p">))</span>
            <span class="n">absolute_sizes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">absolute_size</span><span class="p">)</span>

        <span class="n">absolute_sizes_sum</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">absolute_sizes</span><span class="p">)</span>

        <span class="c1"># if we still need more rows, give them to the first split.</span>
        <span class="c1"># if we have too many rows, remove the extras from the first split that has</span>
        <span class="c1"># enough rows.</span>
        <span class="n">size_difference</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">dataset_size</span> <span class="o">-</span> <span class="n">absolute_sizes_sum</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">size_difference</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">absolute_sizes</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+=</span> <span class="n">size_difference</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">absolute_sizes</span><span class="p">):</span>
                <span class="k">if</span> <span class="n">absolute_sizes</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+</span> <span class="n">size_difference</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">absolute_sizes</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">size_difference</span>
                    <span class="k">break</span>

        <span class="k">if</span> <span class="nb">sum</span><span class="p">(</span><span class="n">absolute_sizes</span><span class="p">)</span> <span class="o">!=</span> <span class="n">dataset_size</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Sum of calculated split sizes </span><span class="si">{}</span><span class="s2"> is not equal to dataset size </span><span class="si">{}</span><span class="s2">.&quot;</span>
                               <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">absolute_sizes_sum</span><span class="p">,</span> <span class="n">dataset_size</span><span class="p">))</span>

        <span class="k">return</span> <span class="n">absolute_sizes</span>

    <span class="nd">@check_split</span>
    <span class="k">def</span> <span class="nf">split</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sizes</span><span class="p">,</span> <span class="n">randomize</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Split the dataset into smaller, non-overlapping datasets.</span>

<span class="sd">        This is a general purpose split function which can be called from any operator in the pipeline.</span>
<span class="sd">        There is another, optimized split function, which will be called automatically if ds.split is</span>
<span class="sd">        called where ds is a MappableDataset.</span>

<span class="sd">        Args:</span>
<span class="sd">            sizes (Union[list[int], list[float]]): If a list of integers [s1, s2, …, sn] is</span>
<span class="sd">                provided, the dataset will be split into n datasets of size s1, size s2, …, size sn</span>
<span class="sd">                respectively. If the sum of all input sizes does not equal the original dataset size, an</span>
<span class="sd">                error will throw.</span>
<span class="sd">                If a list of floats [f1, f2, …, fn] is provided, all floats must be between 0 and 1</span>
<span class="sd">                and must sum to 1, otherwise an error will throw. The dataset will be split into n</span>
<span class="sd">                Datasets of size round(f1*K), round(f2*K), …, round(fn*K) where K is the size of the</span>
<span class="sd">                original dataset.</span>
<span class="sd">                If after rounding:</span>

<span class="sd">                - Any size equals 0, an error will occur.</span>
<span class="sd">                - The sum of split sizes &lt; K, the difference of K - sigma(round(fi * k)) will be added to the first</span>
<span class="sd">                  split.</span>
<span class="sd">                - The sum of split sizes &gt; K, the difference of sigma(round(fi * K)) - K will be removed from the first</span>
<span class="sd">                  large enough split such that it will have at least 1 row after removing the difference.</span>

<span class="sd">            randomize (bool, optional): Determines whether or not to split the data randomly (default=True).</span>
<span class="sd">                If True, the data will be randomly split. Otherwise, each split will be created with</span>
<span class="sd">                consecutive rows from the dataset.</span>

<span class="sd">        Note:</span>
<span class="sd">            1. Dataset cannot be sharded if split is going to be called.</span>
<span class="sd">            2. It is strongly recommended to not shuffle the dataset, but use randomize=True instead.</span>
<span class="sd">               Shuffling the dataset may not be deterministic, which means the data in each split</span>
<span class="sd">               will be different in each epoch.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: If get_dataset_size returns None or is not supported for this dataset.</span>
<span class="sd">            RuntimeError: If `sizes` is list of integers and sum of all elements in sizes does not</span>
<span class="sd">                equal the dataset size.</span>
<span class="sd">            RuntimeError: If `sizes` is list of float and there is a split with size 0 after calculations.</span>
<span class="sd">            RuntimeError: If the dataset is sharded prior to calling split.</span>
<span class="sd">            ValueError: If `sizes` is list of float and not all floats are between 0 and 1, or if the</span>
<span class="sd">                floats don&#39;t sum to 1.</span>

<span class="sd">        Returns:</span>
<span class="sd">            tuple(Dataset), a tuple of datasets that have been split.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # TextFileDataset is not a mappable dataset, so this non-optimized split will be called.</span>
<span class="sd">            &gt;&gt;&gt; # Since many datasets have shuffle on by default, set shuffle to False if split will be called!</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.TextFileDataset(text_file_dataset_dir, shuffle=False)</span>
<span class="sd">            &gt;&gt;&gt; train_dataset, test_dataset = dataset.split([0.9, 0.1])</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_shuffled</span><span class="p">():</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Dataset is shuffled before split.&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_sharded</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Dataset should not be sharded before split.&quot;</span><span class="p">)</span>

        <span class="n">absolute_sizes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_absolute_split_sizes</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span>
        <span class="n">splits</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">rows_to_skip</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">size</span> <span class="ow">in</span> <span class="n">absolute_sizes</span><span class="p">:</span>
            <span class="n">ds</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">randomize</span><span class="p">:</span>
                <span class="c1"># want to shuffle the same way every epoch before split</span>
                <span class="c1"># in alter_tree, shuffle buffer is minimum 10000, so use 10000 here</span>
                <span class="n">ds</span> <span class="o">=</span> <span class="n">ds</span><span class="o">.</span><span class="n">shuffle</span><span class="p">(</span><span class="mi">10000</span><span class="p">)</span>
                <span class="n">ds</span><span class="o">.</span><span class="n">reshuffle_each_epoch</span> <span class="o">=</span> <span class="kc">False</span>

            <span class="k">if</span> <span class="n">rows_to_skip</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">ds</span> <span class="o">=</span> <span class="n">ds</span><span class="o">.</span><span class="n">skip</span><span class="p">(</span><span class="n">rows_to_skip</span><span class="p">)</span>

            <span class="n">ds</span> <span class="o">=</span> <span class="n">ds</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">size</span><span class="p">)</span>
            <span class="n">splits</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">ds</span><span class="p">)</span>

            <span class="n">rows_to_skip</span> <span class="o">+=</span> <span class="n">size</span>

        <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">splits</span><span class="p">)</span>

    <span class="nd">@check_zip_dataset</span>
    <span class="k">def</span> <span class="nf">zip</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">datasets</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Zip the datasets in the sense of input tuple of datasets. Columns in the input datasets must have different</span>
<span class="sd">        name.</span>

<span class="sd">        Args:</span>
<span class="sd">            datasets (Union[tuple, class Dataset]): A tuple of datasets or a single class Dataset</span>
<span class="sd">                to be zipped together with this dataset.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset zipped.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset which is the combination of dataset and dataset_1</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.zip(dataset_1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">datasets</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="n">datasets</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">datasets</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">datasets</span><span class="p">,</span> <span class="n">Dataset</span><span class="p">):</span>
            <span class="n">datasets</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">datasets</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Invalid datasets, expected Dataset object or tuple of Dataset, but got </span><span class="si">%s</span><span class="s2">!&quot;</span> <span class="o">%</span> <span class="n">datasets</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">ZipDataset</span><span class="p">(</span><span class="n">datasets</span><span class="p">)</span>

    <span class="nd">@check_concat</span>
    <span class="k">def</span> <span class="nf">concat</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">datasets</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Concatenate the dataset objects in the input list.</span>
<span class="sd">        Performing &quot;+&quot; operation on dataset objects can achieve the same effect.</span>

<span class="sd">        Note:</span>
<span class="sd">            The column name, and rank and type of the column data must be the same in the input datasets.</span>

<span class="sd">        Args:</span>
<span class="sd">            datasets (Union[list, class Dataset]): A list of datasets or a single class Dataset</span>
<span class="sd">                to be concatenated together with this dataset.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset concatenated.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset by concatenating dataset_1 and dataset_2 with &quot;+&quot; operator</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset_1 + dataset_2</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset by concatenating dataset_1 and dataset_2 with concat operation</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset_1.concat(dataset_2)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">datasets</span><span class="p">,</span> <span class="n">Dataset</span><span class="p">):</span>
            <span class="n">datasets</span> <span class="o">=</span> <span class="p">[</span><span class="bp">self</span><span class="p">]</span> <span class="o">+</span> <span class="p">[</span><span class="n">datasets</span><span class="p">]</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">datasets</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">datasets</span> <span class="o">=</span> <span class="p">[</span><span class="bp">self</span><span class="p">]</span> <span class="o">+</span> <span class="n">datasets</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Invalid datasets, expected Dataset object or list of Dataset, but got </span><span class="si">%s</span><span class="s2">!&quot;</span> <span class="o">%</span> <span class="n">datasets</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">ConcatDataset</span><span class="p">(</span><span class="n">datasets</span><span class="p">)</span>

    <span class="nd">@check_rename</span>
    <span class="k">def</span> <span class="nf">rename</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_columns</span><span class="p">,</span> <span class="n">output_columns</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Rename the columns in input datasets.</span>

<span class="sd">        Args:</span>
<span class="sd">            input_columns (Union[str, list[str]]): List of names of the input columns.</span>
<span class="sd">            output_columns (Union[str, list[str]]): List of names of the output columns.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset renamed.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; input_columns = [&quot;input_col1&quot;, &quot;input_col2&quot;, &quot;input_col3&quot;]</span>
<span class="sd">            &gt;&gt;&gt; output_columns = [&quot;output_col1&quot;, &quot;output_col2&quot;, &quot;output_col3&quot;]</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset where input_col1 is renamed to output_col1, and</span>
<span class="sd">            &gt;&gt;&gt; # input_col2 is renamed to output_col2, and input_col3 is renamed</span>
<span class="sd">            &gt;&gt;&gt; # to output_col3.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.rename(input_columns=input_columns, output_columns=output_columns)</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">return</span> <span class="n">RenameDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_columns</span><span class="p">,</span> <span class="n">output_columns</span><span class="p">)</span>

    <span class="nd">@check_project</span>
    <span class="k">def</span> <span class="nf">project</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Project certain columns in input dataset.</span>

<span class="sd">        The specified columns will be selected from the dataset and passed into</span>
<span class="sd">        the pipeline with the order specified. The other columns are discarded.</span>

<span class="sd">        Args:</span>
<span class="sd">            columns(Union[str, list[str]]): List of names of the columns to project.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset projected.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; columns_to_project = [&quot;column3&quot;, &quot;column1&quot;, &quot;column2&quot;]</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Create a dataset that consists of column3, column1, column2</span>
<span class="sd">            &gt;&gt;&gt; # in that order, regardless of the original order of columns.</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.project(columns=columns_to_project)</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">return</span> <span class="n">ProjectDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">apply</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">apply_func</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Apply a function in this dataset.</span>

<span class="sd">        Args:</span>
<span class="sd">            apply_func (function): A function that must take one &#39;Dataset&#39; as an argument and</span>
<span class="sd">                                   return a preprocessed &#39;Dataset&#39;.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset applied by the function.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Declare an apply_func function which returns a Dataset object</span>
<span class="sd">            &gt;&gt;&gt; def apply_func(data):</span>
<span class="sd">            ...     data = data.batch(2)</span>
<span class="sd">            ...     return data</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Use apply to call apply_func</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.apply(apply_func)</span>

<span class="sd">        Raises:</span>
<span class="sd">            TypeError: If apply_func is not a function.</span>
<span class="sd">            TypeError: If apply_func doesn&#39;t return a Dataset.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">apply_func</span><span class="p">,</span> <span class="s1">&#39;__call__&#39;</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;apply_func must be a function.&quot;</span><span class="p">)</span>

        <span class="n">dataset</span> <span class="o">=</span> <span class="n">apply_func</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">Dataset</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;apply_func must return a dataset.&quot;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">dataset</span>

    <span class="nd">@check_device_send</span>
    <span class="k">def</span> <span class="nf">device_que</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">send_epoch_end</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">create_data_info_queue</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return a transferred Dataset that transfers data through a device.</span>

<span class="sd">        Args:</span>
<span class="sd">            send_epoch_end (bool, optional): Whether to send end of sequence to device or not (default=True).</span>
<span class="sd">            create_data_info_queue (bool, optional): Whether to create queue which stores</span>
<span class="sd">                types and shapes of data or not(default=False).</span>

<span class="sd">        Note:</span>
<span class="sd">            If device is Ascend, features of data will be transferred one by one. The limitation</span>
<span class="sd">            of data transmission per time is 256M.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Dataset, dataset for transferring.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">to_device</span><span class="p">(</span><span class="n">send_epoch_end</span><span class="o">=</span><span class="n">send_epoch_end</span><span class="p">,</span> <span class="n">create_data_info_queue</span><span class="o">=</span><span class="n">create_data_info_queue</span><span class="p">)</span>

    <span class="nd">@check_device_send</span>
    <span class="k">def</span> <span class="nf">to_device</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">send_epoch_end</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">create_data_info_queue</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Transfer data from CPU to GPU or Ascend or other devices.</span>

<span class="sd">        Args:</span>
<span class="sd">            send_epoch_end (bool, optional): Whether to send the end of sequence to device or not (default=True).</span>
<span class="sd">            create_data_info_queue (bool, optional): Whether to create queue which stores</span>
<span class="sd">                types and shapes of data or not(default=False).</span>

<span class="sd">        Note:</span>
<span class="sd">            If device is Ascend, features of data will be transferred one by one. The limitation</span>
<span class="sd">            of data transmission per second is 256M.</span>

<span class="sd">        Returns:</span>
<span class="sd">            TransferDataset, dataset for transferring.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: If distribution file path is given but failed to read.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">TransferDataset</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">send_epoch_end</span><span class="p">,</span> <span class="n">create_data_info_queue</span><span class="p">)</span>

    <span class="nd">@check_save</span>
    <span class="k">def</span> <span class="nf">save</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">file_name</span><span class="p">,</span> <span class="n">num_files</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">file_type</span><span class="o">=</span><span class="s1">&#39;mindrecord&#39;</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Save the dynamic data processed by the dataset pipeline in common dataset format.</span>
<span class="sd">        Supported dataset formats: &#39;mindrecord&#39; only</span>

<span class="sd">        Implicit type casting exists when saving data as &#39;mindrecord&#39;. The transform table shows how to do type casting.</span>

<span class="sd">        .. list-table:: Implicit Type Casting when Saving as &#39;mindrecord&#39;</span>
<span class="sd">           :widths: 25 25 50</span>
<span class="sd">           :header-rows: 1</span>

<span class="sd">           * - Type in &#39;dataset&#39;</span>
<span class="sd">             - Type in &#39;mindrecord&#39;</span>
<span class="sd">             - Details</span>
<span class="sd">           * - bool</span>
<span class="sd">             - None</span>
<span class="sd">             - Not supported</span>
<span class="sd">           * - int8</span>
<span class="sd">             - int32</span>
<span class="sd">             -</span>
<span class="sd">           * - uint8</span>
<span class="sd">             - bytes(1D uint8)</span>
<span class="sd">             - Drop dimension</span>
<span class="sd">           * - int16</span>
<span class="sd">             - int32</span>
<span class="sd">             -</span>
<span class="sd">           * - uint16</span>
<span class="sd">             - int32</span>
<span class="sd">             -</span>
<span class="sd">           * - int32</span>
<span class="sd">             - int32</span>
<span class="sd">             -</span>
<span class="sd">           * - uint32</span>
<span class="sd">             - int64</span>
<span class="sd">             -</span>
<span class="sd">           * - int64</span>
<span class="sd">             - int64</span>
<span class="sd">             -</span>
<span class="sd">           * - uint64</span>
<span class="sd">             - None</span>
<span class="sd">             - Not supported</span>
<span class="sd">           * - float16</span>
<span class="sd">             - float32</span>
<span class="sd">             -</span>
<span class="sd">           * - float32</span>
<span class="sd">             - float32</span>
<span class="sd">             -</span>
<span class="sd">           * - float64</span>
<span class="sd">             - float64</span>
<span class="sd">             -</span>
<span class="sd">           * - string</span>
<span class="sd">             - string</span>
<span class="sd">             - Multi-dimensional string not supported</span>

<span class="sd">        Note:</span>
<span class="sd">            1. To save the samples in order, set dataset&#39;s shuffle to False and num_files to 1.</span>
<span class="sd">            2. Before calling the function, do not use batch operator, repeat operator or data augmentation operators</span>
<span class="sd">               with random attribute in map operator.</span>
<span class="sd">            3. When array dimension is variable, one-dimensional arrays or</span>
<span class="sd">               multi-dimensional arrays with variable dimension 0 are supported.</span>
<span class="sd">            4. Mindrecord does not support uint64, multi-dimensional uint8(drop dimension) nor</span>
<span class="sd">               multi-dimensional string.</span>

<span class="sd">        Args:</span>
<span class="sd">            file_name (str): Path to dataset file.</span>
<span class="sd">            num_files (int, optional): Number of dataset files (default=1).</span>
<span class="sd">            file_type (str, optional): Dataset format (default=&#39;mindrecord&#39;).</span>

<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ir_tree</span><span class="p">,</span> <span class="n">api_tree</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>

        <span class="n">runtime_context</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonRuntimeContext</span><span class="p">()</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">Init</span><span class="p">()</span>
        <span class="n">consumer</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonSaveToDisk</span><span class="p">(</span><span class="n">file_name</span><span class="p">,</span> <span class="n">num_files</span><span class="p">,</span> <span class="n">file_type</span><span class="p">)</span>
        <span class="n">consumer</span><span class="o">.</span><span class="n">Init</span><span class="p">(</span><span class="n">ir_tree</span><span class="p">)</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">AssignConsumer</span><span class="p">(</span><span class="n">consumer</span><span class="p">)</span>

        <span class="n">consumer</span><span class="o">.</span><span class="n">Save</span><span class="p">()</span>
        <span class="n">_set_dataset_permissions</span><span class="p">(</span><span class="n">file_name</span><span class="p">,</span> <span class="n">num_files</span><span class="p">)</span>
        <span class="k">del</span> <span class="n">api_tree</span>

    <span class="nd">@check_tuple_iterator</span>
    <span class="k">def</span> <span class="nf">create_tuple_iterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_epochs</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">output_numpy</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">do_copy</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Create an iterator over the dataset. The datatype retrieved back will be a list of ndarrays.</span>

<span class="sd">        To specify which columns to list and the order needed, use columns_list. If columns_list</span>
<span class="sd">        is not provided, the order of the columns will remain unchanged.</span>

<span class="sd">        Args:</span>
<span class="sd">            columns (list[str], optional): List of columns to be used to specify the order of columns</span>
<span class="sd">                (default=None, means all columns).</span>
<span class="sd">            num_epochs (int, optional): Maximum number of epochs that iterator can be iterated.</span>
<span class="sd">                (default=-1, iterator can be iterated infinite number of epochs)</span>
<span class="sd">            output_numpy (bool, optional): Whether or not to output NumPy datatype.</span>
<span class="sd">                If output_numpy=False, iterator will output MSTensor (default=False).</span>
<span class="sd">            do_copy (bool, optional): when output data type is mindspore.Tensor,</span>
<span class="sd">                use this param to select the conversion method, only take False for better performance (default=True).</span>

<span class="sd">        Returns:</span>
<span class="sd">            Iterator, tuple iterator over the dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; iterator = dataset.create_tuple_iterator()</span>
<span class="sd">            &gt;&gt;&gt; for item in iterator:</span>
<span class="sd">            ...     # item is a list</span>
<span class="sd">            ...     print(type(item))</span>
<span class="sd">            ...     break</span>
<span class="sd">            &lt;class &#39;list&#39;&gt;</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">output_numpy</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">output_numpy</span> <span class="o">=</span> <span class="kc">False</span>

        <span class="k">if</span> <span class="n">Dataset</span><span class="o">.</span><span class="n">_noop_mode</span><span class="p">():</span>
            <span class="k">return</span> <span class="n">DummyIterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;tuple&#39;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">TupleIterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">,</span> <span class="n">num_epochs</span><span class="p">,</span> <span class="n">output_numpy</span><span class="p">,</span> <span class="n">do_copy</span><span class="p">)</span>

    <span class="nd">@check_dict_iterator</span>
    <span class="k">def</span> <span class="nf">create_dict_iterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_epochs</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">output_numpy</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Create an iterator over the dataset. The data retrieved will be a dictionary datatype.</span>

<span class="sd">        The order of the columns in the dictionary may not be the same as the original order.</span>

<span class="sd">        Args:</span>
<span class="sd">            num_epochs (int, optional): Maximum number of epochs that iterator can be iterated</span>
<span class="sd">                (default=-1, iterator can be iterated infinite number of epochs).</span>
<span class="sd">            output_numpy (bool, optional): Whether or not to output NumPy datatype,</span>
<span class="sd">                if output_numpy=False, iterator will output MSTensor (default=False).</span>

<span class="sd">        Returns:</span>
<span class="sd">            Iterator, dictionary iterator over the dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; iterator = dataset.create_dict_iterator()</span>
<span class="sd">            &gt;&gt;&gt; for item in iterator:</span>
<span class="sd">            ...     # item is a dict</span>
<span class="sd">            ...     print(type(item))</span>
<span class="sd">            ...     break</span>
<span class="sd">            &lt;class &#39;dict&#39;&gt;</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">output_numpy</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">output_numpy</span> <span class="o">=</span> <span class="kc">False</span>

        <span class="k">if</span> <span class="n">Dataset</span><span class="o">.</span><span class="n">_noop_mode</span><span class="p">():</span>
            <span class="k">return</span> <span class="n">DummyIterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;dict&#39;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">DictIterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_epochs</span><span class="p">,</span> <span class="n">output_numpy</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__iter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Create an iterator over the dataset.&quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_tuple_iterator</span><span class="p">(</span><span class="n">num_epochs</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">input_indexs</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get Input Index Information</span>

<span class="sd">        Returns:</span>
<span class="sd">            tuple, tuple of the input index information.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; # set input_indexs</span>
<span class="sd">            &gt;&gt;&gt; dataset.input_indexs = 10</span>
<span class="sd">            &gt;&gt;&gt; print(dataset.input_indexs)</span>
<span class="sd">            10</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_input_indexs</span> <span class="o">!=</span> <span class="p">():</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_input_indexs</span>

        <span class="c1"># find input_indexes of children</span>
        <span class="n">children_input_index</span> <span class="o">=</span> <span class="p">[</span><span class="n">child</span><span class="o">.</span><span class="n">input_indexs</span> <span class="k">for</span> <span class="n">child</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">]</span>

        <span class="c1"># in case of more than one child, return the first input_indexes</span>
        <span class="k">for</span> <span class="n">cix</span> <span class="ow">in</span> <span class="n">children_input_index</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">cix</span> <span class="o">!=</span> <span class="p">():</span>
                <span class="k">return</span> <span class="n">cix</span>

        <span class="c1"># if all children&#39;s input_indexes are () or the node is a leaf</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_input_indexs</span>

    <span class="nd">@input_indexs</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">input_indexs</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_input_indexs</span> <span class="o">=</span> <span class="n">value</span>

    <span class="k">def</span> <span class="nf">copy_batch_size</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span> <span class="o">=</span> <span class="n">value</span>

    <span class="k">def</span> <span class="nf">_init_tree_getters</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get pipeline information.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ir_tree</span><span class="p">,</span> <span class="n">api_tree</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>

        <span class="n">runtime_context</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonRuntimeContext</span><span class="p">()</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">Init</span><span class="p">()</span>
        <span class="n">getter</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">TreeGetters</span><span class="p">()</span>
        <span class="n">getter</span><span class="o">.</span><span class="n">Init</span><span class="p">(</span><span class="n">ir_tree</span><span class="p">)</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">AssignConsumer</span><span class="p">(</span><span class="n">getter</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">getter</span><span class="p">,</span> <span class="n">runtime_context</span><span class="p">,</span> <span class="n">api_tree</span>

    <span class="k">def</span> <span class="nf">__init_size_getter</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get pipeline information.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ir_tree</span><span class="p">,</span> <span class="n">api_tree</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>

        <span class="n">runtime_context</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonRuntimeContext</span><span class="p">()</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">Init</span><span class="p">()</span>
        <span class="n">getter</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">DatasetSizeGetters</span><span class="p">()</span>
        <span class="n">getter</span><span class="o">.</span><span class="n">Init</span><span class="p">(</span><span class="n">ir_tree</span><span class="p">)</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">AssignConsumer</span><span class="p">(</span><span class="n">getter</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">getter</span><span class="p">,</span> <span class="n">runtime_context</span><span class="p">,</span> <span class="n">api_tree</span>

    <span class="k">def</span> <span class="nf">get_col_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the names of the columns in dataset.</span>

<span class="sd">        Returns:</span>
<span class="sd">            list, list of column names in the dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; col_names = dataset.get_col_names()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_col_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_tree_getters</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_col_names</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetColumnNames</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">close_pool</span><span class="p">()</span>
            <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">notify_watchdog</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_col_names</span>

    <span class="k">def</span> <span class="nf">output_shapes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get the shapes of output data.</span>

<span class="sd">        Returns:</span>
<span class="sd">            list, list of shapes of each column.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; output_shapes = dataset.output_shapes()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_tree_getters</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetOutputShapes</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_types</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetOutputTypes</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">close_pool</span><span class="p">()</span>
            <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">notify_watchdog</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span><span class="p">[</span><span class="mi">0</span><span class="p">]:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_dynamic_output_shapes</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span>

    <span class="k">def</span> <span class="nf">output_types</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get the types of output data.</span>

<span class="sd">        Returns:</span>
<span class="sd">            list, list of data types.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; output_types = dataset.output_types()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_types</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_tree_getters</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetOutputShapes</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_types</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetOutputTypes</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">close_pool</span><span class="p">()</span>
            <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">notify_watchdog</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span><span class="p">[</span><span class="mi">0</span><span class="p">]:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_dynamic_output_shapes</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_types</span>

    <span class="k">def</span> <span class="nf">get_dataset_size</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the number of batches in an epoch.</span>

<span class="sd">        Returns:</span>
<span class="sd">            int, number of batches.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; dataset_size = dataset.get_dataset_size()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">__init_size_getter</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetDatasetSize</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">close_pool</span><span class="p">()</span>
            <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">notify_watchdog</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span>

    <span class="k">def</span> <span class="nf">set_dynamic_columns</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set dynamic shape information of source data, it should be set after the pipeline is defined.</span>

<span class="sd">        Args:</span>
<span class="sd">            columns (dict): A dict contains shape information of each column in dataset.</span>
<span class="sd">                The value of shape[i] is :py:obj:`None` indicates that the data length of shape[i] is dynamic.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; import numpy as np</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; def generator1():</span>
<span class="sd">            ...     for i in range(1, 100):</span>
<span class="sd">            ...         yield np.ones((16, i, 83)), np.array(i)</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.GeneratorDataset(generator1, [&quot;data1&quot;, &quot;data2&quot;])</span>
<span class="sd">            &gt;&gt;&gt; dataset.set_dynamic_columns(columns={&quot;data1&quot;: [16, None, 83], &quot;data2&quot;: []})</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">columns</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Pass a dict to set dynamic shape, example: {</span><span class="se">\&quot;</span><span class="s2">data1</span><span class="se">\&quot;</span><span class="s2">: [16, None, 256]}&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">columns</span>

    <span class="k">def</span> <span class="nf">dynamic_min_max_shapes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get minimum and maximum data length of dynamic source data, for dynamic graph compilation.</span>

<span class="sd">        Returns:</span>
<span class="sd">            lists, min_shapes, max_shapes of source data.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; import numpy as np</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; def generator1():</span>
<span class="sd">            ...     for i in range(1, 100):</span>
<span class="sd">            ...         yield np.ones((16, i, 83)), np.array(i)</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.GeneratorDataset(generator1, [&quot;data1&quot;, &quot;data2&quot;])</span>
<span class="sd">            &gt;&gt;&gt; dataset.set_dynamic_columns(columns={&quot;data1&quot;: [16, None, 83], &quot;data2&quot;: []})</span>
<span class="sd">            &gt;&gt;&gt; min_shapes, max_shapes = dataset.dynamic_min_max_shapes()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_dynamic_output_shapes</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">__check_dynamic_column_name</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">,</span> <span class="n">dataset_columns</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">column</span> <span class="ow">in</span> <span class="n">dynamic_columns</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">column</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">dataset_columns</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;dynamic column [&quot;</span> <span class="o">+</span> <span class="n">column</span> <span class="o">+</span> <span class="s2">&quot;] does not match any column in dataset: &quot;</span> <span class="o">+</span>
                                   <span class="nb">str</span><span class="p">(</span><span class="n">dataset_columns</span><span class="p">))</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">__check_dynamic_column_shape</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">col</span><span class="p">,</span> <span class="n">dynamic_columns</span><span class="p">):</span>
        <span class="n">shape_mismatch</span> <span class="o">=</span> <span class="s2">&quot;dynamic column [&quot;</span> <span class="o">+</span> <span class="n">col</span> <span class="o">+</span> <span class="s2">&quot;] with shape &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">])</span> <span class="o">+</span> \
                         <span class="s2">&quot; does not match dataset column [&quot;</span> <span class="o">+</span> <span class="n">col</span> <span class="o">+</span> <span class="s2">&quot;] with shape &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">data</span><span class="p">[</span><span class="n">col</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">data</span><span class="p">[</span><span class="n">col</span><span class="p">]</span><span class="o">.</span><span class="n">ndim</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">]):</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="n">shape_mismatch</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">dim</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">])):</span>
            <span class="k">if</span> <span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">][</span><span class="n">dim</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">][</span><span class="n">dim</span><span class="p">]</span> <span class="o">!=</span> <span class="n">data</span><span class="p">[</span><span class="n">col</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="n">shape_mismatch</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_dynamic_output_shapes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get dynamic information of source data.</span>

<span class="sd">        Returns:</span>
<span class="sd">            lists, dynamic_shapes, min_shapes, max_shapes of source data.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span><span class="p">[</span><span class="mi">1</span><span class="p">]:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;dynamic_columns is not set, call set_dynamic_columns() by final Dataset Op.&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> \
                <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_output_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_min_shapes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">saved_max_shapes</span>

        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Calculating dynamic shape of input data, this will take a few minutes...&quot;</span><span class="p">)</span>
        <span class="c1"># Assume data1 shape is dynamic, data2 shape is fix</span>
        <span class="c1"># {&quot;data1&quot;: [batch_size, None, feat_len], &quot;data2&quot;: [batch_size, feat_len]}</span>
        <span class="n">dynamic_columns</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_setting</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
        <span class="c1"># [&quot;data1&quot;, &quot;data2&quot;]</span>
        <span class="n">dataset_columns</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_col_names</span><span class="p">()</span>
        <span class="n">Dataset</span><span class="o">.</span><span class="n">__check_dynamic_column_name</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">,</span> <span class="n">dataset_columns</span><span class="p">)</span>

        <span class="c1"># Shape[1] of data1 is variable</span>
        <span class="c1"># {&quot;data1&quot;: {(batch_size, 100, feat_len), (16, 200, 83)}, &quot;data2&quot;: {(batch_size, feat_len)}}</span>
        <span class="n">column_shape_set</span> <span class="o">=</span> <span class="p">{</span><span class="n">col</span><span class="p">:</span> <span class="nb">set</span><span class="p">()</span> <span class="k">for</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">dataset_columns</span><span class="p">}</span>
        <span class="n">dataset_size_counter</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">data</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_dict_iterator</span><span class="p">(</span><span class="n">num_epochs</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">output_numpy</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
            <span class="n">dataset_size_counter</span> <span class="o">+=</span> <span class="mi">1</span>
            <span class="k">for</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">data</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                <span class="k">if</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">dynamic_columns</span><span class="p">:</span>
                    <span class="n">Dataset</span><span class="o">.</span><span class="n">__check_dynamic_column_shape</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">col</span><span class="p">,</span> <span class="n">dynamic_columns</span><span class="p">)</span>
                <span class="n">column_shape_set</span><span class="p">[</span><span class="n">col</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">data</span><span class="p">[</span><span class="n">col</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span>

        <span class="c1"># we get dataset_size after dryrun</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="n">dataset_size_counter</span>

        <span class="n">min_shapes</span><span class="p">,</span> <span class="n">max_shapes</span><span class="p">,</span> <span class="n">dynamic_shapes</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(),</span> <span class="nb">list</span><span class="p">(),</span> <span class="nb">list</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">col</span><span class="p">,</span> <span class="n">shape_set</span> <span class="ow">in</span> <span class="n">column_shape_set</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape_set</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">col</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">dynamic_columns</span><span class="p">:</span>
                    <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;column [&quot;</span> <span class="o">+</span> <span class="n">col</span> <span class="o">+</span> <span class="s2">&quot;] has dynamic shape but not set by set_dynamic_columns()&quot;</span> <span class="o">+</span>
                                       <span class="s2">&quot;, shapes of [&quot;</span> <span class="o">+</span> <span class="n">col</span> <span class="o">+</span> <span class="s2">&quot;]: &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">shape_set</span><span class="p">)))</span>
                <span class="n">shape_npy</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">shape_set</span><span class="p">))</span>
                <span class="n">max_shape</span> <span class="o">=</span> <span class="n">shape_npy</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
                <span class="n">min_shape</span> <span class="o">=</span> <span class="n">shape_npy</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>

                <span class="c1"># Set min shape to 1 due to unknown shuffle</span>
                <span class="n">min_shape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">],</span> <span class="kc">None</span><span class="p">),</span> <span class="mi">1</span><span class="p">,</span> <span class="n">min_shape</span><span class="p">)</span>
                <span class="c1"># Set dynamic dim to -1 for ME</span>
                <span class="n">dynamic_shape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">],</span> <span class="kc">None</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">])</span>

                <span class="n">max_shapes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">max_shape</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span>
                <span class="n">min_shapes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">min_shape</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span>
                <span class="n">dynamic_shapes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dynamic_shape</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="c1"># Also append fix shape to keep order of column shape</span>
                <span class="n">fix_shape</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">shape_set</span><span class="p">)[</span><span class="mi">0</span><span class="p">])</span>
                <span class="n">max_shapes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">fix_shape</span><span class="p">)</span>
                <span class="n">min_shapes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">fix_shape</span><span class="p">)</span>
                <span class="n">dynamic_shapes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">fix_shape</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">dynamic_columns</span><span class="p">:</span>
                    <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;column [&quot;</span> <span class="o">+</span> <span class="n">col</span> <span class="o">+</span> <span class="s2">&quot;] has no dynamic shape but set by set_dynamic_columns()&quot;</span><span class="p">)</span>
                    <span class="c1"># Set min shape to 1 due to unknown shuffle</span>
                    <span class="n">min_shapes</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">],</span> <span class="kc">None</span><span class="p">),</span> <span class="mi">1</span><span class="p">,</span> <span class="n">fix_shape</span><span class="p">)</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span>
                    <span class="c1"># Set dynamic dim to -1 for ME</span>
                    <span class="n">dynamic_shapes</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">dynamic_columns</span><span class="p">[</span><span class="n">col</span><span class="p">],</span> <span class="kc">None</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">fix_shape</span><span class="p">)</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">dynamic_shapes</span><span class="p">,</span> <span class="n">min_shapes</span><span class="p">,</span> <span class="n">max_shapes</span>

    <span class="k">def</span> <span class="nf">num_classes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get the number of classes in a dataset.</span>

<span class="sd">        Returns:</span>
<span class="sd">            int, number of classes.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; num_classes = dataset.num_classes()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_num_classes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_tree_getters</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_num_classes</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetNumClasses</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">close_pool</span><span class="p">()</span>
            <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">notify_watchdog</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_num_classes</span> <span class="o">==</span> <span class="o">-</span><span class="mi">1</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">None</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_num_classes</span>

    <span class="k">def</span> <span class="nf">get_sync_notifiers</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">get_sync_notifiers</span><span class="p">()</span>
        <span class="k">return</span> <span class="p">{}</span>

    <span class="k">def</span> <span class="nf">disable_sync</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">disable_sync</span><span class="p">()</span>
        <span class="k">return</span> <span class="p">{}</span>

    <span class="k">def</span> <span class="nf">is_sync</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">is_sync</span><span class="p">()</span>
        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">sync_update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">condition_name</span><span class="p">,</span> <span class="n">num_batch</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">data</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Release a blocking condition and trigger callback with given data.</span>

<span class="sd">        Args:</span>
<span class="sd">            condition_name (str): The condition name that is used to toggle sending next row.</span>
<span class="sd">            num_batch (Union[int, None]): The number of batches (rows) that are released.</span>
<span class="sd">                When num_batch is None, it will default to the number specified by the</span>
<span class="sd">                sync_wait operator (default=None).</span>
<span class="sd">            data (Any): The data passed to the callback, user defined (default=None).</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="p">(</span><span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">num_batch</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">and</span> <span class="n">num_batch</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span> <span class="ow">or</span> \
                <span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">num_batch</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">and</span> <span class="n">num_batch</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">):</span>
            <span class="c1"># throwing exception, disable all sync_wait in pipeline</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">disable_sync</span><span class="p">()</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Sync_update batch size can only be positive integer, got : </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">num_batch</span><span class="p">))</span>
        <span class="n">notifiers_dict</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_sync_notifiers</span><span class="p">()</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">condition_name</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Argument condition_name with value </span><span class="si">{}</span><span class="s2"> is not of type str, but got </span><span class="si">{}</span><span class="s2">.&quot;</span>
                            <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">condition_name</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="n">condition_name</span><span class="p">)))</span>
        <span class="k">if</span> <span class="n">condition_name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">notifiers_dict</span><span class="p">:</span>
            <span class="c1"># throwing exception, disable all sync_wait in pipeline</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">disable_sync</span><span class="p">()</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Condition name not found.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">num_batch</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">num_batch</span> <span class="o">*=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_batch_size</span><span class="p">()</span>
        <span class="n">notifiers_dict</span><span class="p">[</span><span class="n">condition_name</span><span class="p">](</span><span class="n">num_batch</span><span class="p">,</span> <span class="n">data</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">get_batch_size</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the size of batch.</span>

<span class="sd">        Returns:</span>
<span class="sd">            int, the number of data in a batch.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; batch_size = dataset.get_batch_size()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_tree_getters</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetBatchSize</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_batch_size</span>

    <span class="k">def</span> <span class="nf">get_repeat_count</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get the replication times in RepeatDataset (default is 1).</span>

<span class="sd">        Returns:</span>
<span class="sd">            int, the count of repeat.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; repeat_count = dataset.get_repeat_count()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_repeat_count</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">runtime_getter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_tree_getters</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_repeat_count</span> <span class="o">=</span> <span class="n">runtime_getter</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">GetRepeatCount</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_repeat_count</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_repeat_count</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_repeat_count</span>

    <span class="k">def</span> <span class="nf">get_class_indexing</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the class index.</span>

<span class="sd">        Returns:</span>
<span class="sd">            dict, a str-to-int mapping from label name to index.</span>
<span class="sd">            dict, a str-to-list&lt;int&gt; mapping from label name to index for Coco ONLY. The second number</span>
<span class="sd">            in the list is used to indicate the super category.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; class_indexing = dataset.get_class_indexing()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">get_class_indexing</span><span class="p">()</span>
        <span class="k">return</span> <span class="p">{}</span>

    <span class="k">def</span> <span class="nf">reset</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Reset the dataset for next epoch.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">is_shuffled</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Returns True if the dataset or its children is shuffled.&quot;&quot;&quot;</span>
        <span class="k">for</span> <span class="n">input_dataset</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">input_dataset</span><span class="o">.</span><span class="n">is_shuffled</span><span class="p">():</span>
                <span class="k">return</span> <span class="kc">True</span>

        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">is_sharded</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Returns True if the dataset or its children is sharded.&quot;&quot;&quot;</span>
        <span class="k">for</span> <span class="n">input_dataset</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">input_dataset</span><span class="o">.</span><span class="n">is_sharded</span><span class="p">():</span>
                <span class="k">return</span> <span class="kc">True</span>

        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">post_parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ir_node</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache</span><span class="p">:</span>
            <span class="n">ir_node</span> <span class="o">=</span> <span class="n">ir_node</span><span class="o">.</span><span class="n">set_cache_client</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cache</span><span class="o">.</span><span class="n">cache_client</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span><span class="p">:</span>
            <span class="n">ir_node</span> <span class="o">=</span> <span class="n">ir_node</span><span class="o">.</span><span class="n">set_num_workers</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">ir_node</span>


<span class="k">class</span> <span class="nc">VisionBaseDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a vision source dataset which produces content to the data pipeline.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">children</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">TextBaseDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a text source dataset which produces content to the data pipeline.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">children</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">build_vocab</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">,</span> <span class="n">freq_range</span><span class="p">,</span> <span class="n">top_k</span><span class="p">,</span> <span class="n">special_tokens</span><span class="p">,</span> <span class="n">special_first</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Function to create a Vocab from source dataset.</span>
<span class="sd">        Desired source dataset is a text type dataset.</span>

<span class="sd">        Build a vocab from a dataset. This would collect all the unique words in a dataset and return a vocab</span>
<span class="sd">        which contains top_k most frequent words (if top_k is specified)</span>

<span class="sd">        Args:</span>

<span class="sd">            columns(Union[str, list[str]]): Column names to get words from.</span>
<span class="sd">            freq_range(tuple[int]): A tuple of integers (min_frequency, max_frequency). Words within the frequency</span>
<span class="sd">                range will be stored.</span>
<span class="sd">                Naturally 0 &lt;= min_frequency &lt;= max_frequency &lt;= total_words. min_frequency/max_frequency</span>
<span class="sd">                can be set to default, which corresponds to 0/total_words separately.</span>
<span class="sd">            top_k(int): Number of words to be built into vocab. top_k most frequent words are</span>
<span class="sd">                taken. The top_k is taken after freq_range. If not enough top_k, all words will be taken</span>
<span class="sd">            special_tokens(list[str]): A list of strings, each one is a special token.</span>
<span class="sd">            special_first(bool): Whether special_tokens will be prepended/appended to vocab, If special_tokens</span>
<span class="sd">                is specified and special_first is set to default, special_tokens will be prepended.</span>

<span class="sd">        Returns:</span>
<span class="sd">            Vocab, vocab built from the dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; import numpy as np</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; def gen_corpus():</span>
<span class="sd">            ...     # key: word, value: number of occurrences, reason for using letters is so their order is apparent</span>
<span class="sd">            ...     corpus = {&quot;Z&quot;: 4, &quot;Y&quot;: 4, &quot;X&quot;: 4, &quot;W&quot;: 3, &quot;U&quot;: 3, &quot;V&quot;: 2, &quot;T&quot;: 1}</span>
<span class="sd">            ...     for k, v in corpus.items():</span>
<span class="sd">            ...         yield (np.array([k] * v, dtype=&#39;S&#39;),)</span>
<span class="sd">            &gt;&gt;&gt; column_names = [&quot;column1&quot;]</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.GeneratorDataset(gen_corpus, column_names)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.build_vocab(columns=[&quot;column1&quot;],</span>
<span class="sd">            ...                               freq_range=(1, 10), top_k=5,</span>
<span class="sd">            ...                               special_tokens=[&quot;&lt;pad&gt;&quot;, &quot;&lt;unk&gt;&quot;],</span>
<span class="sd">            ...                               special_first=True)</span>

<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">vocab</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">Vocab</span><span class="p">()</span>
        <span class="n">columns</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">columns</span><span class="p">,</span> <span class="p">[])</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">columns</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">columns</span> <span class="o">=</span> <span class="p">[</span><span class="n">columns</span><span class="p">]</span>

        <span class="n">freq_range</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">freq_range</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">9223372036854775807</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">freq_range</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">freq_range</span> <span class="o">=</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">freq_range</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
        <span class="k">if</span> <span class="n">freq_range</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">freq_range</span> <span class="o">=</span> <span class="p">(</span><span class="n">freq_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">9223372036854775807</span><span class="p">)</span>
        <span class="n">special_tokens</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">special_tokens</span><span class="p">,</span> <span class="p">[])</span>
        <span class="n">top_k</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">top_k</span><span class="p">,</span> <span class="mi">9223372036854775807</span><span class="p">)</span>

        <span class="n">ir_tree</span><span class="p">,</span> <span class="n">api_tree</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>

        <span class="c1"># vocab node</span>
        <span class="n">vocab_node</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">BuildVocabNode</span><span class="p">(</span><span class="n">ir_tree</span><span class="p">,</span> <span class="n">vocab</span><span class="p">,</span> <span class="n">columns</span><span class="p">,</span> <span class="n">freq_range</span><span class="p">,</span> <span class="n">top_k</span><span class="p">,</span> <span class="n">special_tokens</span><span class="p">,</span> <span class="n">special_first</span><span class="p">)</span>

        <span class="n">runtime_context</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonRuntimeContext</span><span class="p">()</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">Init</span><span class="p">()</span>

        <span class="c1"># build vocab</span>
        <span class="n">consumer</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonBuildVocabConsumer</span><span class="p">()</span>
        <span class="n">consumer</span><span class="o">.</span><span class="n">Init</span><span class="p">(</span><span class="n">vocab_node</span><span class="p">)</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">AssignConsumer</span><span class="p">(</span><span class="n">consumer</span><span class="p">)</span>

        <span class="n">consumer</span><span class="o">.</span><span class="n">Start</span><span class="p">()</span>
        <span class="k">del</span> <span class="n">api_tree</span>

        <span class="k">return</span> <span class="n">vocab</span>

    <span class="k">def</span> <span class="nf">build_sentencepiece_vocab</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="n">character_coverage</span><span class="p">,</span> <span class="n">model_type</span><span class="p">,</span> <span class="n">params</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Function to create a SentencePieceVocab from source dataset.</span>
<span class="sd">        Desired source dataset is a text type dataset.</span>

<span class="sd">        Args:</span>

<span class="sd">            columns(list[str]): Column names to get words from.</span>
<span class="sd">            vocab_size(int): Vocabulary size.</span>
<span class="sd">            character_coverage(int): Percentage of characters covered by the model, must be between</span>
<span class="sd">                0.98 and 1.0 Good defaults are: 0.9995 for languages with rich character sets like</span>
<span class="sd">                Japanese or Chinese character sets, and 1.0 for other languages with small character sets</span>
<span class="sd">                like English or Latin.</span>
<span class="sd">            model_type(SentencePieceModel): Model type. Choose from unigram (default), bpe, char, or word.</span>
<span class="sd">                The input sentence must be pretokenized when using word type.</span>
<span class="sd">            params(dict): Any extra optional parameters of sentencepiece library according to your raw data</span>

<span class="sd">        Returns:</span>
<span class="sd">            SentencePieceVocab, vocab built from the dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; from mindspore.dataset.text import SentencePieceModel</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # You can construct any text dataset as source, take TextFileDataset as example.</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.TextFileDataset(&quot;/path/to/sentence/piece/vocab/file&quot;, shuffle=False)</span>
<span class="sd">            &gt;&gt;&gt; dataset = dataset.build_sentencepiece_vocab([&quot;text&quot;], 5000, 0.9995, SentencePieceModel.UNIGRAM, {})</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model_type</span><span class="p">,</span> <span class="n">SentencePieceModel</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Argument model_type with value </span><span class="si">{0}</span><span class="s2"> is not of type SentencePieceModel, but got </span><span class="si">{1}</span><span class="s2">.&quot;</span>\
                            <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">model_type</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="n">model_type</span><span class="p">)))</span>
        <span class="n">model_type</span> <span class="o">=</span> <span class="n">DE_C_INTER_SENTENCEPIECE_MODE</span><span class="p">[</span><span class="n">model_type</span><span class="p">]</span>
        <span class="n">vocab</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">SentencePieceVocab</span><span class="p">()</span>

        <span class="n">ir_tree</span><span class="p">,</span> <span class="n">api_tree</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>

        <span class="c1"># vocab node</span>
        <span class="n">vocab_node</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">BuildSentenceVocabNode</span><span class="p">(</span><span class="n">ir_tree</span><span class="p">,</span> <span class="n">vocab</span><span class="p">,</span> <span class="n">columns</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">,</span> <span class="n">character_coverage</span><span class="p">,</span> <span class="n">model_type</span><span class="p">,</span>
                                                <span class="n">params</span><span class="p">)</span>

        <span class="n">runtime_context</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonRuntimeContext</span><span class="p">()</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">Init</span><span class="p">()</span>

        <span class="c1"># build vocab</span>
        <span class="n">consumer</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonBuildVocabConsumer</span><span class="p">()</span>
        <span class="n">consumer</span><span class="o">.</span><span class="n">Init</span><span class="p">(</span><span class="n">vocab_node</span><span class="p">)</span>
        <span class="n">runtime_context</span><span class="o">.</span><span class="n">AssignConsumer</span><span class="p">(</span><span class="n">consumer</span><span class="p">)</span>

        <span class="n">consumer</span><span class="o">.</span><span class="n">Start</span><span class="p">()</span>
        <span class="k">del</span> <span class="n">api_tree</span>

        <span class="k">return</span> <span class="n">vocab</span>


<span class="k">class</span> <span class="nc">AudioBaseDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a audio source dataset which produces content to the data pipeline.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">children</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">UnionBaseDataset</span><span class="p">(</span><span class="n">VisionBaseDataset</span><span class="p">,</span> <span class="n">TextBaseDataset</span><span class="p">,</span> <span class="n">AudioBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a union source dataset which produces content to the data pipeline.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">children</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SourceDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a source dataset which produces content to the data pipeline.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_samples</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">shuffle</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">num_shards</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">shard_id</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_samples</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">num_samples</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_shards</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">num_shards</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shard_id</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">shard_id</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">shuffle</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shuffle</span><span class="p">,</span> <span class="p">(</span><span class="nb">bool</span><span class="p">,</span> <span class="n">Shuffle</span><span class="p">)):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;shuffle must be of boolean or enum of &#39;Shuffle&#39; values like &#39;Shuffle.GLOBAL&#39; or &quot;</span>
                            <span class="s2">&quot;&#39;Shuffle.FILES&#39; or &#39;Shuffle.INFILE&#39;.&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="mi">2</span>  <span class="c1"># Global shuffle</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shuffle</span><span class="p">,</span> <span class="n">Shuffle</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">shuffle</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">shuffle</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="mi">2</span>  <span class="c1"># Global shuffle</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="mi">0</span>  <span class="c1"># No shuffle</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">shuffle</span> <span class="o">==</span> <span class="n">Shuffle</span><span class="o">.</span><span class="n">GLOBAL</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="mi">2</span>  <span class="c1"># Global shuffle</span>
            <span class="k">elif</span> <span class="n">shuffle</span> <span class="o">==</span> <span class="n">Shuffle</span><span class="o">.</span><span class="n">FILES</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="mi">1</span>  <span class="c1"># Files shuffle</span>
            <span class="k">elif</span> <span class="n">shuffle</span> <span class="o">==</span> <span class="n">Shuffle</span><span class="o">.</span><span class="n">INFILE</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="mi">3</span>  <span class="c1"># Infile shuffle</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_find_files</span><span class="p">(</span><span class="n">patterns</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Utility function to search for files with the given glob patterns.</span>

<span class="sd">        Args:</span>
<span class="sd">            patterns (Union[str, list[str]]): String or list of patterns to be searched.</span>

<span class="sd">        Returns:</span>
<span class="sd">            list, list of files.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">patterns</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">patterns</span> <span class="o">=</span> <span class="p">[</span><span class="n">patterns</span><span class="p">]</span>

        <span class="n">file_list</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">unmatched_patterns</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">pattern</span> <span class="ow">in</span> <span class="n">patterns</span><span class="p">:</span>
            <span class="n">matches</span> <span class="o">=</span> <span class="p">[</span><span class="n">match</span> <span class="k">for</span> <span class="n">match</span> <span class="ow">in</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span><span class="n">pattern</span><span class="p">,</span> <span class="n">recursive</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">isfile</span><span class="p">(</span><span class="n">match</span><span class="p">)]</span>

            <span class="k">if</span> <span class="n">matches</span><span class="p">:</span>
                <span class="n">file_list</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">matches</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">unmatched_patterns</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">pattern</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">unmatched_patterns</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The following patterns did not match any files: </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">unmatched_patterns</span><span class="p">))</span>

        <span class="k">if</span> <span class="n">file_list</span><span class="p">:</span>  <span class="c1"># not empty</span>
            <span class="k">return</span> <span class="n">file_list</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The list of path names matching the patterns is empty.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">is_shuffled</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">&gt;</span> <span class="mi">0</span>

    <span class="k">def</span> <span class="nf">is_sharded</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_shards</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_shards</span> <span class="o">&gt;</span> <span class="mi">1</span>
        <span class="k">return</span> <span class="kc">False</span>


<span class="k">class</span> <span class="nc">MappableDataset</span><span class="p">(</span><span class="n">SourceDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class to represent a source dataset which supports use of samplers.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Dataset has to implement parse method.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">sampler</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_samples</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">shuffle</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_shards</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">shard_id</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">num_samples</span><span class="o">=</span><span class="n">num_samples</span><span class="p">,</span> <span class="n">shuffle</span><span class="o">=</span><span class="n">shuffle</span><span class="p">,</span>
                         <span class="n">num_shards</span><span class="o">=</span><span class="n">num_shards</span><span class="p">,</span> <span class="n">shard_id</span><span class="o">=</span><span class="n">shard_id</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shuffle_flag</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">shuffle</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sampler</span> <span class="o">=</span> <span class="n">samplers</span><span class="o">.</span><span class="n">select_sampler</span><span class="p">(</span><span class="n">num_samples</span><span class="p">,</span> <span class="n">sampler</span><span class="p">,</span> <span class="n">shuffle</span><span class="p">,</span> <span class="n">num_shards</span><span class="p">,</span> <span class="n">shard_id</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">add_sampler</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">new_sampler</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Add a sampler for current dataset.</span>

<span class="sd">        Args:</span>
<span class="sd">            new_sampler (Sampler): The sampler to be added as the parent sampler for current dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; # use a DistributedSampler instead</span>
<span class="sd">            &gt;&gt;&gt; new_sampler = ds.DistributedSampler(10, 2)</span>
<span class="sd">            &gt;&gt;&gt; dataset.add_sampler(new_sampler)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># note: By adding a sampler, the sampled IDs will flow to new_sampler</span>
        <span class="c1"># after first passing through the current samplers attached to this dataset.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">new_sampler</span><span class="o">.</span><span class="n">add_child</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">sampler</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sampler</span> <span class="o">=</span> <span class="n">new_sampler</span>

    <span class="k">def</span> <span class="nf">use_sampler</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">new_sampler</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Make the current dataset use the new_sampler provided by other API.</span>

<span class="sd">        Args:</span>
<span class="sd">            new_sampler (Sampler): The sampler to use for the current dataset.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # dataset is an instance object of Dataset</span>
<span class="sd">            &gt;&gt;&gt; # use a DistributedSampler instead</span>
<span class="sd">            &gt;&gt;&gt; new_sampler = ds.DistributedSampler(10, 2)</span>
<span class="sd">            &gt;&gt;&gt; dataset.use_sampler(new_sampler)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">new_sampler</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Input sampler can not be None.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">new_sampler</span><span class="p">,</span> <span class="p">(</span><span class="n">samplers</span><span class="o">.</span><span class="n">BuiltinSampler</span><span class="p">,</span> <span class="n">samplers</span><span class="o">.</span><span class="n">Sampler</span><span class="p">)):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Input sampler is not an instance of a sampler.&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">sampler</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">sampler</span><span class="o">.</span><span class="n">child_sampler</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_sampler</span><span class="p">(</span><span class="n">new_sampler</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">is_shuffled</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">sampler</span><span class="o">.</span><span class="n">is_shuffled</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">is_sharded</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">sampler</span><span class="o">.</span><span class="n">is_sharded</span><span class="p">()</span>

    <span class="nd">@check_split</span>
    <span class="k">def</span> <span class="nf">split</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sizes</span><span class="p">,</span> <span class="n">randomize</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Split the dataset into smaller, non-overlapping datasets.</span>

<span class="sd">        Args:</span>
<span class="sd">            sizes (Union[list[int], list[float]]): If a list of integers [s1, s2, …, sn] is</span>
<span class="sd">                provided, the dataset will be split into n datasets of size s1, size s2, …, size sn</span>
<span class="sd">                respectively. If the sum of all sizes does not equal the original dataset size, an</span>
<span class="sd">                error will occur.</span>
<span class="sd">                If a list of floats [f1, f2, …, fn] is provided, all floats must be between 0 and 1</span>
<span class="sd">                and must sum to 1, otherwise an error will occur. The dataset will be split into n</span>
<span class="sd">                Datasets of size round(f1*K), round(f2*K), …, round(fn*K) where K is the size of the</span>
<span class="sd">                original dataset.</span>
<span class="sd">                If after rounding:</span>

<span class="sd">                - Any size equals 0, an error will occur.</span>
<span class="sd">                - The sum of split sizes &lt; K, the difference will be added to the first split.</span>
<span class="sd">                - The sum of split sizes &gt; K, the difference will be removed from the first large</span>
<span class="sd">                  enough split such that it will have at least 1 row after removing the difference.</span>

<span class="sd">            randomize (bool, optional): Determines whether or not to split the data randomly (default=True).</span>
<span class="sd">                If True, the data will be randomly split. Otherwise, each split will be created with</span>
<span class="sd">                consecutive rows from the dataset.</span>

<span class="sd">        Note:</span>
<span class="sd">            1. There is an optimized split function, which will be called automatically when the dataset</span>
<span class="sd">               that calls this function is a MappableDataset.</span>
<span class="sd">            2. Dataset should not be sharded if split is going to be called. Instead, create a</span>
<span class="sd">               DistributedSampler and specify a split to shard after splitting. If the dataset is</span>
<span class="sd">               sharded after a split, it is strongly recommended setting the same seed in each instance</span>
<span class="sd">               of execution, otherwise each shard may not be part of the same split (see Examples).</span>
<span class="sd">            3. It is strongly recommended to not shuffle the dataset, but use randomize=True instead.</span>
<span class="sd">               Shuffling the dataset may not be deterministic, which means the data in each split</span>
<span class="sd">               will be different in each epoch. Furthermore, if sharding occurs after split, each</span>
<span class="sd">               shard may not be part of the same split.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: If get_dataset_size returns None or is not supported for this dataset.</span>
<span class="sd">            RuntimeError: If `sizes` is list of integers and sum of all elements in sizes does not</span>
<span class="sd">                equal the dataset size.</span>
<span class="sd">            RuntimeError: If `sizes` is list of float and there is a split with size 0 after calculations.</span>
<span class="sd">            RuntimeError: If the dataset is sharded prior to calling split.</span>
<span class="sd">            ValueError: If `sizes` is list of float and not all floats are between 0 and 1, or if the</span>
<span class="sd">                floats don&#39;t sum to 1.</span>

<span class="sd">        Returns:</span>
<span class="sd">            tuple(Dataset), a tuple of datasets that have been split.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; # Since many datasets have shuffle on by default, set shuffle to False if split will be called!</span>
<span class="sd">            &gt;&gt;&gt; dataset = ds.ImageFolderDataset(image_folder_dataset_dir, shuffle=False)</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # Set the seed, and tell split to use this seed when randomizing.</span>
<span class="sd">            &gt;&gt;&gt; # This is needed because sharding will be done later</span>
<span class="sd">            &gt;&gt;&gt; ds.config.set_seed(58)</span>
<span class="sd">            &gt;&gt;&gt; train_dataset, test_dataset = dataset.split([0.9, 0.1])</span>
<span class="sd">            &gt;&gt;&gt;</span>
<span class="sd">            &gt;&gt;&gt; # To shard the train dataset, use a DistributedSampler</span>
<span class="sd">            &gt;&gt;&gt; train_sampler = ds.DistributedSampler(10, 2)</span>
<span class="sd">            &gt;&gt;&gt; train_dataset.use_sampler(train_sampler)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_shuffled</span><span class="p">():</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Dataset is shuffled before split.&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_sharded</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Dataset should not be sharded before split.&quot;</span><span class="p">)</span>

        <span class="n">absolute_sizes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_absolute_split_sizes</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span>
        <span class="n">splits</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">current_split_start_index</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">size</span> <span class="ow">in</span> <span class="n">absolute_sizes</span><span class="p">:</span>
            <span class="n">ds</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
            <span class="n">ds</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">if</span> <span class="n">randomize</span><span class="p">:</span>
                <span class="c1"># want to shuffle the same way every epoch before split, we are assuming</span>
                <span class="c1"># that the user will call set_seed</span>
                <span class="n">random_sampler</span> <span class="o">=</span> <span class="n">samplers</span><span class="o">.</span><span class="n">RandomSampler</span><span class="p">()</span>
                <span class="n">random_sampler</span><span class="o">.</span><span class="n">reshuffle_each_epoch</span> <span class="o">=</span> <span class="kc">False</span>
                <span class="n">ds</span><span class="o">.</span><span class="n">add_sampler</span><span class="p">(</span><span class="n">random_sampler</span><span class="p">)</span>

            <span class="n">subset_sampler</span> <span class="o">=</span> <span class="n">samplers</span><span class="o">.</span><span class="n">SequentialSampler</span><span class="p">(</span><span class="n">current_split_start_index</span><span class="p">,</span> <span class="n">size</span><span class="p">)</span>
            <span class="n">ds</span><span class="o">.</span><span class="n">add_sampler</span><span class="p">(</span><span class="n">subset_sampler</span><span class="p">)</span>

            <span class="c1"># add sequential sampler, so that if user calls use_sampler, we will</span>
            <span class="c1"># get rid of the sequential sampler instead of something we need</span>
            <span class="n">ds</span><span class="o">.</span><span class="n">add_sampler</span><span class="p">(</span><span class="n">samplers</span><span class="o">.</span><span class="n">SequentialSampler</span><span class="p">())</span>

            <span class="n">splits</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">ds</span><span class="p">)</span>

            <span class="n">current_split_start_index</span> <span class="o">+=</span> <span class="n">size</span>

        <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">splits</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">BucketBatchByLengthDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying BucketBatchByLength operator to the input dataset.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">column_names</span><span class="p">,</span> <span class="n">bucket_boundaries</span><span class="p">,</span> <span class="n">bucket_batch_sizes</span><span class="p">,</span> <span class="n">element_length_function</span><span class="p">,</span>
                 <span class="n">pad_info</span><span class="p">,</span> <span class="n">pad_to_bucket_boundary</span><span class="p">,</span> <span class="n">drop_remainder</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">column_names</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">column_names</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bucket_boundaries</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">bucket_boundaries</span><span class="p">,</span> <span class="p">[])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bucket_batch_sizes</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">bucket_batch_sizes</span><span class="p">,</span> <span class="p">[])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">element_length_function</span> <span class="o">=</span> <span class="n">element_length_function</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pad_info</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">pad_info</span><span class="p">,</span> <span class="p">{})</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pad_to_bucket_boundary</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">pad_to_bucket_boundary</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">drop_remainder</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">drop_remainder</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">BucketBatchByLengthNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">column_names</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">bucket_boundaries</span><span class="p">,</span>
                                           <span class="bp">self</span><span class="o">.</span><span class="n">bucket_batch_sizes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">element_length_function</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">pad_info</span><span class="p">,</span>
                                           <span class="bp">self</span><span class="o">.</span><span class="n">pad_to_bucket_boundary</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">drop_remainder</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_check_shm_usage</span><span class="p">(</span><span class="n">num_worker</span><span class="p">,</span> <span class="n">queue_size</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="p">,</span> <span class="n">num_queues</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Check sufficient shared memory is available for shared memory queues</span>
<span class="sd">    when training in parallel mode.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">threshold_ratio</span> <span class="o">=</span> <span class="mf">0.8</span>
    <span class="k">if</span> <span class="n">platform</span><span class="o">.</span><span class="n">system</span><span class="p">()</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">{</span><span class="s2">&quot;windows&quot;</span><span class="p">,</span> <span class="s2">&quot;darwin&quot;</span><span class="p">}:</span>
        <span class="n">device_num</span> <span class="o">=</span> <span class="n">_get_device_num</span><span class="p">()</span>
        <span class="c1"># In the cluster, _get_device_num indicates the number of the entire cluster. The maximum number of cards</span>
        <span class="c1"># on the ascend server is 8.</span>
        <span class="k">if</span> <span class="n">device_num</span> <span class="o">&gt;</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;device_target&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;Ascend&quot;</span><span class="p">:</span>
            <span class="n">device_num</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">device_num</span><span class="p">,</span> <span class="mi">8</span><span class="p">)</span>
        <span class="n">shm_estimate_usage</span> <span class="o">=</span> <span class="n">device_num</span> <span class="o">*</span> <span class="n">num_worker</span> <span class="o">*</span> <span class="n">num_queues</span> <span class="o">*</span> \
                             <span class="p">(</span><span class="n">queue_size</span> <span class="o">+</span> <span class="mi">2</span><span class="p">)</span> <span class="o">*</span> <span class="n">max_rowsize</span> <span class="o">*</span> <span class="mi">1024</span> <span class="o">*</span> <span class="mi">1024</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">shm_available</span> <span class="o">=</span> <span class="n">psutil</span><span class="o">.</span><span class="n">disk_usage</span><span class="p">(</span><span class="s1">&#39;/dev/shm&#39;</span><span class="p">)</span><span class="o">.</span><span class="n">free</span>
            <span class="k">if</span> <span class="n">shm_estimate_usage</span> <span class="o">&gt;=</span> <span class="n">threshold_ratio</span> <span class="o">*</span> <span class="n">shm_available</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                    <span class="s2">&quot;Insufficient shared memory available. Required: </span><span class="si">{}</span><span class="s2">, Available: </span><span class="si">{}</span><span class="s2">. &quot;</span>
                    <span class="s2">&quot;The required memory can&#39;t exceed 80</span><span class="si">% o</span><span class="s2">f the available shared memory, &quot;</span>
                    <span class="s2">&quot;it&#39;s recommended to reduce memory usage by following methods:</span><span class="se">\n</span><span class="s2">&quot;</span>
                    <span class="s2">&quot;1. reduce value of parameter max_rowsize or num_parallel_workers.</span><span class="se">\n</span><span class="s2">&quot;</span>
                    <span class="s2">&quot;2. reduce prefetch size by set_prefetch_size().</span><span class="se">\n</span><span class="s2">&quot;</span>
                    <span class="s2">&quot;3. disable shared memory by set_enable_shared_mem().&quot;</span>
                    <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">shm_estimate_usage</span><span class="p">,</span> <span class="n">shm_available</span><span class="p">))</span>
        <span class="k">except</span> <span class="ne">FileNotFoundError</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Expected /dev/shm to exist.&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">BatchDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Batch operator to the input dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be batched.</span>
<span class="sd">        batch_size (Union[int, function]): The number of rows each batch is created with. An</span>
<span class="sd">            int or callable which takes exactly 1 parameter, BatchInfo.</span>
<span class="sd">        drop_remainder (bool, optional): Determines whether or not to drop the last</span>
<span class="sd">            possibly incomplete batch (default=False). If True, and if there are less</span>
<span class="sd">            than batch_size rows available to make the last batch, then those rows will</span>
<span class="sd">            be dropped and not propagated to the child node.</span>
<span class="sd">        num_parallel_workers (int, optional): Number of workers to process the dataset in parallel (default=None).</span>
<span class="sd">        per_batch_map (callable, optional): Per batch map callable. A callable which takes</span>
<span class="sd">            (list[Tensor], list[Tensor], ..., BatchInfo) as input parameters. Each list[Tensor] represents a batch of</span>
<span class="sd">            Tensors on a given column. The number of lists should match with number of entries in input_columns. The</span>
<span class="sd">            last parameter of the callable must always be a BatchInfo object.</span>
<span class="sd">        input_columns (Union[str, list[str]], optional): List of names of the input columns. The size of the list must</span>
<span class="sd">            match with signature of per_batch_map callable.</span>
<span class="sd">        output_columns (Union[str, list[str]], optional): List of names assigned to the columns outputted by</span>
<span class="sd">            the last operation. This parameter is mandatory if len(input_columns) !=</span>
<span class="sd">            len(output_columns). The size of this list must match the number of output</span>
<span class="sd">            columns of the last operation. (default=None, output columns will have the same</span>
<span class="sd">            name as the input columns, i.e., the columns will be replaced).</span>
<span class="sd">        column_order (Union[str, list[str]], optional): Specifies the list of all the columns you need in the whole</span>
<span class="sd">                dataset. The parameter is required when len(input_column) != len(output_column). Caution: the list here</span>
<span class="sd">                is not just the columns specified in parameter input_columns and output_columns.</span>
<span class="sd">        pad_info (dict, optional): Whether to perform padding on selected columns. pad_info={&quot;col1&quot;:([224,224],0)}</span>
<span class="sd">            will pad column with name &quot;col1&quot; to a tensor of size [224,224] and fill the missing with 0.</span>
<span class="sd">        max_rowsize(int, optional): Maximum size of row in MB that is used for shared memory allocation to copy</span>
<span class="sd">            data between processes.  This is only used if python_multiprocessing is set to True (default=16).</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">drop_remainder</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">per_batch_map</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">input_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">output_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">column_order</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">pad_info</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">python_multiprocessing</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="mi">16</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">BatchDataset</span><span class="o">.</span><span class="n">_is_ancestor_of_repeat</span><span class="p">(</span><span class="n">input_dataset</span><span class="p">):</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Repeat is located before batch, data from two epochs can be batched together.&quot;</span><span class="p">)</span>

        <span class="n">BatchDataset</span><span class="o">.</span><span class="n">_update_batch_size_for_syncwait</span><span class="p">(</span><span class="n">input_dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">)</span>

        <span class="c1"># if batch_size is callable, set batch_size to 1 and batch_size_func to that callable function</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span> <span class="o">=</span> <span class="n">batch_size</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">callable</span><span class="p">(</span><span class="n">batch_size</span><span class="p">)</span> <span class="k">else</span> <span class="mi">1</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_func</span> <span class="o">=</span> <span class="kc">None</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">callable</span><span class="p">(</span><span class="n">batch_size</span><span class="p">)</span> <span class="k">else</span> <span class="n">batch_size</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">drop_remainder</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">drop_remainder</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span> <span class="o">=</span> <span class="n">per_batch_map</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">input_columns</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">output_columns</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">column_order</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">column_order</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">pad</span> <span class="o">=</span> <span class="nb">bool</span><span class="p">(</span><span class="n">pad_info</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pad_info</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">pad_info</span><span class="p">,</span> <span class="nb">dict</span><span class="p">())</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">python_multiprocessing</span> <span class="o">=</span> <span class="n">python_multiprocessing</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">hook</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">workers</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span> <span class="o">=</span> <span class="n">max_rowsize</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">BatchNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">drop_remainder</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">pad</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span><span class="p">,</span>
                             <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">column_order</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_func</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span><span class="p">,</span>
                             <span class="bp">self</span><span class="o">.</span><span class="n">pad_info</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_is_ancestor_of_repeat</span><span class="p">(</span><span class="n">dataset</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Utility function to find the case where repeat is used before batch.</span>

<span class="sd">        Args:</span>
<span class="sd">             dataset (Dataset): Dataset to be checked.</span>

<span class="sd">        Returns:</span>
<span class="sd">            bool, whether repeat is used before batch.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">RepeatDataset</span><span class="p">):</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="n">flag</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">for</span> <span class="n">input_dataset</span> <span class="ow">in</span> <span class="n">dataset</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="n">flag</span> <span class="o">=</span> <span class="n">flag</span> <span class="o">|</span> <span class="n">BatchDataset</span><span class="o">.</span><span class="n">_is_ancestor_of_repeat</span><span class="p">(</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">flag</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_update_batch_size_for_syncwait</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Utility function to notify batch size to sync_wait.</span>

<span class="sd">        Args:</span>
<span class="sd">             dataset (Dataset): Dataset to be checked.</span>
<span class="sd">             batch_size (int): batch size to notify.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">SyncWaitDataset</span><span class="p">):</span>
            <span class="n">dataset</span><span class="o">.</span><span class="n">update_sync_batch_size</span><span class="p">(</span><span class="n">batch_size</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">input_dataset</span> <span class="ow">in</span> <span class="n">dataset</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="n">BatchDataset</span><span class="o">.</span><span class="n">_update_batch_size_for_syncwait</span><span class="p">(</span><span class="n">input_dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__deepcopy__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">memodict</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">__safe_deepcopy__</span><span class="p">(</span><span class="n">memodict</span><span class="p">,</span> <span class="n">exclude</span><span class="o">=</span><span class="p">(</span><span class="s2">&quot;per_batch_map&quot;</span><span class="p">,</span> <span class="s2">&quot;batch_size_func&quot;</span><span class="p">,</span> <span class="s2">&quot;__transfer_dataset__&quot;</span><span class="p">))</span>

    <span class="c1"># Iterator bootstrap will be called on iterator construction.</span>
    <span class="c1"># A deep copy of Dataset object is created prior of iterator_bootstrap.</span>
    <span class="c1"># This method will create per iterator process pool and bind pyfunc execution to the pool.</span>
    <span class="k">def</span> <span class="nf">iterator_bootstrap</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Per iterator bootstrap callback.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">python_multiprocessing</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;per_batch_map is None so python_multiprocessing does not work.&quot;</span><span class="p">)</span>
                <span class="k">return</span>
            <span class="n">arg_q_list</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">res_q_list</span> <span class="o">=</span> <span class="p">[]</span>

            <span class="c1"># If user didn&#39;t specify num_parallel_workers, set it to default</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">num_parallel</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">num_parallel</span> <span class="o">=</span> <span class="n">get_num_parallel_workers</span><span class="p">()</span>

            <span class="k">if</span> <span class="n">get_enable_shared_mem</span><span class="p">():</span>
                <span class="n">_check_shm_usage</span><span class="p">(</span><span class="n">num_parallel</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
                <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_parallel</span><span class="p">):</span>
                    <span class="n">arg_q_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_SharedQueue</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span><span class="p">))</span>
                    <span class="n">res_q_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_SharedQueue</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span><span class="p">))</span>

            <span class="c1"># Construct pool with the callable list</span>
            <span class="c1"># The callable list and _pyfunc_worker_init are used to pass lambda function in to subprocesses</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="o">=</span> <span class="n">multiprocessing</span><span class="o">.</span><span class="n">Pool</span><span class="p">(</span><span class="n">processes</span><span class="o">=</span><span class="n">num_parallel</span><span class="p">,</span>
                                                     <span class="n">initializer</span><span class="o">=</span><span class="n">_pyfunc_worker_init</span><span class="p">,</span>
                                                     <span class="n">initargs</span><span class="o">=</span><span class="p">([</span><span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span><span class="p">],</span> <span class="n">arg_q_list</span><span class="p">,</span> <span class="n">res_q_list</span><span class="p">))</span>

            <span class="n">idx</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="k">global</span> <span class="n">_OP_NAME</span><span class="p">,</span> <span class="n">_OP_PROCESS</span><span class="p">,</span> <span class="n">_LOCK</span>
            <span class="n">op_id</span> <span class="o">=</span> <span class="n">_OP_NAME</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="p">)]</span>
            <span class="n">process_id</span> <span class="o">=</span> <span class="p">{</span><span class="n">op_id</span><span class="p">:</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="nb">set</span><span class="p">()]}</span>
            <span class="c1"># obtain process id from multiprocessing.pool</span>
            <span class="k">for</span> <span class="n">pool</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">_pool</span><span class="p">:</span>  <span class="c1"># pylint: disable=W0212</span>
                <span class="n">process_id</span><span class="p">[</span><span class="n">op_id</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">pool</span><span class="o">.</span><span class="n">pid</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">workers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">pool</span><span class="p">)</span>
            <span class="k">with</span> <span class="n">_LOCK</span><span class="p">:</span>
                <span class="n">_OP_PROCESS</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">process_id</span><span class="p">)</span>

            <span class="c1"># Wrap per_batch_map into _PythonCallable</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span> <span class="o">=</span> <span class="n">_PythonCallable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span><span class="p">,</span> <span class="n">idx</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="p">,</span> <span class="n">arg_q_list</span><span class="p">,</span> <span class="n">res_q_list</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">hook</span> <span class="o">=</span> <span class="n">_ExceptHookHandler</span><span class="p">()</span>

            <span class="c1"># batch will launch a watch dog thread to monitoring sub processes</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_launch_watch_dog</span><span class="p">()</span>

            <span class="n">atexit</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="n">_mp_pool_exit_preprocess</span><span class="p">)</span>
            <span class="c1"># If Python version greater than 3.8, we need to close ThreadPool in atexit for unclean pool teardown.</span>
            <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">version_info</span> <span class="o">&gt;=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">8</span><span class="p">):</span>
                <span class="n">atexit</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">close</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span> <span class="o">=</span> <span class="n">FuncWrapper</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">per_batch_map</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_launch_watch_dog</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">platform</span><span class="o">.</span><span class="n">system</span><span class="p">()</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="o">!=</span> <span class="s1">&#39;windows&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Event</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Thread</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">_watch_dog</span><span class="p">,</span> <span class="n">args</span><span class="o">=</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">eot</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">workers</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span><span class="o">.</span><span class="n">daemon</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_abort_watchdog</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">eot</span><span class="o">.</span><span class="n">is_set</span><span class="p">():</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">eot</span><span class="o">.</span><span class="n">set</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;process_pool&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;watch_dog&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;eot&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_abort_watchdog</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">BatchInfo</span><span class="p">(</span><span class="n">cde</span><span class="o">.</span><span class="n">CBatchInfo</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Only the batch size function and per_batch_map of the batch operator can dynamically adjust parameters</span>
<span class="sd">    based on the number of batches and epochs during training.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">get_batch_num</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the batch number of the current batch.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span>

    <span class="k">def</span> <span class="nf">get_epoch_num</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the epoch number of the current batch.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span>


<span class="k">class</span> <span class="nc">BlockReleasePair</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The blocking condition class used by SyncWaitDataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        init_release_rows (int): Number of lines to allow through the pipeline.</span>
<span class="sd">        callback (function): The callback function that will be called when release is called (default=None).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">init_release_rows</span><span class="p">,</span> <span class="n">callback</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">init_release_rows</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">and</span> <span class="n">init_release_rows</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;release_rows need to be greater than 0.&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">row_count</span> <span class="o">=</span> <span class="o">-</span><span class="n">init_release_rows</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cv</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Condition</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">callback</span> <span class="o">=</span> <span class="n">callback</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">default_rows</span> <span class="o">=</span> <span class="n">init_release_rows</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">disable</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">__deepcopy__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">memodict</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="nf">reset</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">row_count</span> <span class="o">=</span> <span class="o">-</span><span class="bp">self</span><span class="o">.</span><span class="n">default_rows</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="o">.</span><span class="n">notify_all</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">update_batched_size</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">):</span>
        <span class="c1"># sanity check</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">and</span> <span class="n">batch_size</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;batch_size need to be greater than 0.&quot;</span><span class="p">)</span>

        <span class="c1"># should only use before the pipeline creates</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">row_count</span> <span class="o">*=</span> <span class="n">batch_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">default_rows</span> <span class="o">*=</span> <span class="n">batch_size</span>

    <span class="k">def</span> <span class="nf">block_func</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Function for handing blocking condition.</span>

<span class="sd">        Returns:</span>
<span class="sd">            bool, True.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="p">:</span>
            <span class="c1"># if disable is true, the always evaluate to true</span>
            <span class="n">not_time_out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="o">.</span><span class="n">wait_for</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">row_count</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">disable</span><span class="p">),</span>
                                            <span class="n">timeout</span><span class="o">=</span><span class="n">get_callback_timeout</span><span class="p">())</span>
            <span class="c1"># time_out will be False if time out occurs</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">not_time_out</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Timeout happened in sync_wait, maybe dataset.sync_update(condition=...) &quot;</span>
                               <span class="s2">&quot;is not added after dataset.create_dict_iterator(...), now disabling lock.&quot;</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">disable</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">row_count</span> <span class="o">+=</span> <span class="mi">1</span>
        <span class="k">return</span> <span class="kc">True</span>

    <span class="k">def</span> <span class="nf">release_func</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pass_rows</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">data</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">pass_rows</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">pass_rows</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">default_rows</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">row_count</span> <span class="o">-=</span> <span class="n">pass_rows</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">callback</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">callback</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="o">.</span><span class="n">notify_all</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">disable_lock</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">disable</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cv</span><span class="o">.</span><span class="n">notify_all</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">SyncWaitDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of adding a blocking condition to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input dataset to apply flow control.</span>
<span class="sd">        num_batch (int): Number of batches without blocking at the start of each epoch.</span>
<span class="sd">        condition_name (str): Condition name that is used to toggle sending next row.</span>
<span class="sd">        callback (function): Callback function that will be invoked when sync_update is called (default=None).</span>

<span class="sd">    Raises:</span>
<span class="sd">        RuntimeError: If condition name already exists.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">condition_name</span><span class="p">,</span> <span class="n">num_batch</span><span class="p">,</span> <span class="n">callback</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>

        <span class="c1"># set to the default value, waiting for the batch to update it</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_condition_name</span> <span class="o">=</span> <span class="n">condition_name</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">num_batch</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">and</span> <span class="n">num_batch</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;num_batch need to be greater than 0.&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_pair</span> <span class="o">=</span> <span class="n">BlockReleasePair</span><span class="p">(</span><span class="n">num_batch</span><span class="p">,</span> <span class="n">callback</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_condition_name</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">get_sync_notifiers</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Condition name is already in use.&quot;</span><span class="p">)</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Please remember to add dataset.sync_update(condition=</span><span class="si">%s</span><span class="s2">), otherwise hanging will result. &quot;</span>
                    <span class="s2">&quot;If dataset.sync_update(condition=</span><span class="si">%s</span><span class="s2">) has already been added, you can ignore the info.&quot;</span><span class="p">,</span>
                    <span class="n">condition_name</span><span class="p">,</span> <span class="n">condition_name</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">SyncWaitNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">_condition_name</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pair</span><span class="o">.</span><span class="n">block_func</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">get_sync_notifiers</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="p">{</span><span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">get_sync_notifiers</span><span class="p">(),</span> <span class="o">**</span><span class="p">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_condition_name</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pair</span><span class="o">.</span><span class="n">release_func</span><span class="p">}}</span>

    <span class="k">def</span> <span class="nf">is_sync</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="kc">True</span>

    <span class="k">def</span> <span class="nf">update_sync_batch_size</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">and</span> <span class="n">batch_size</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;num_batch need to be greater than 0.&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_pair</span><span class="o">.</span><span class="n">update_batched_size</span><span class="p">(</span><span class="n">batch_size</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">disable_sync</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Disabling Sync&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_pair</span><span class="o">.</span><span class="n">disable_lock</span><span class="p">()</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_is_ancestor_of_batch</span><span class="p">(</span><span class="n">dataset</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Utility function to find the case where sync_wait is used before batch.</span>

<span class="sd">        Args:</span>
<span class="sd">             dataset (Dataset): Dataset to be checked.</span>

<span class="sd">        Returns:</span>
<span class="sd">            bool, whether sync_wait is used before batch.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">BatchDataset</span><span class="p">):</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="n">flag</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">for</span> <span class="n">input_dataset</span> <span class="ow">in</span> <span class="n">dataset</span><span class="o">.</span><span class="n">children</span><span class="p">:</span>
            <span class="n">flag</span> <span class="o">=</span> <span class="n">flag</span> <span class="o">|</span> <span class="n">SyncWaitDataset</span><span class="o">.</span><span class="n">_is_ancestor_of_batch</span><span class="p">(</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">flag</span>

    <span class="k">def</span> <span class="nf">iterator_bootstrap</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_pair</span><span class="o">.</span><span class="n">reset</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">ShuffleDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Shuffle operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be shuffled.</span>
<span class="sd">        buffer_size (int): Size of the buffer.</span>

<span class="sd">    Raises:</span>
<span class="sd">        RuntimeError: If exist sync operators before shuffle.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">buffer_size</span> <span class="o">=</span> <span class="n">buffer_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reshuffle_each_epoch</span> <span class="o">=</span> <span class="kc">True</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_sync</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;No shuffle after sync operators.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">ShuffleNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">buffer_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshuffle_each_epoch</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">is_shuffled</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="kc">True</span>


<span class="c1"># This wait function is for cleaning zombie subprocesses</span>
<span class="k">def</span> <span class="nf">wait_pid</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This function is used by the main process to release subprocess resources.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="k">while</span> <span class="kc">True</span><span class="p">:</span>
            <span class="n">child_pid</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">waitpid</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">os</span><span class="o">.</span><span class="n">WNOHANG</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">child_pid</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="k">break</span>
    <span class="k">except</span> <span class="ne">OSError</span><span class="p">:</span>
        <span class="c1"># waitpid may be failed for some reasons so we ignore this error</span>
        <span class="k">pass</span>


<span class="c1"># Terminate subprocess launched by multiprocessing.pool</span>
<span class="k">def</span> <span class="nf">_terminate_process</span><span class="p">(</span><span class="n">workers</span><span class="p">):</span>
    <span class="k">for</span> <span class="n">w</span> <span class="ow">in</span> <span class="n">workers</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">w</span><span class="o">.</span><span class="n">exitcode</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">w</span><span class="o">.</span><span class="n">terminate</span><span class="p">()</span>
    <span class="k">for</span> <span class="n">w</span> <span class="ow">in</span> <span class="n">workers</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">w</span><span class="o">.</span><span class="n">_closed</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>  <span class="c1"># pylint: disable=W0212</span>
            <span class="n">w</span><span class="o">.</span><span class="n">join</span><span class="p">()</span>


<span class="c1"># Monitor the exit number of subprocesses</span>
<span class="k">def</span> <span class="nf">_monitor_subprocess_exit</span><span class="p">(</span><span class="n">workers</span><span class="p">):</span>
    <span class="n">subprocess_exit_num</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">for</span> <span class="n">w</span> <span class="ow">in</span> <span class="n">workers</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">w</span><span class="o">.</span><span class="n">exitcode</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">subprocess_exit_num</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">return</span> <span class="n">subprocess_exit_num</span>


<span class="c1"># Dataset need _watch_dog thread to monitoring fork multi-processing,</span>
<span class="c1"># and thread can&#39;t be a member function otherwise python won&#39;t collect and release resources.</span>
<span class="k">def</span> <span class="nf">_watch_dog</span><span class="p">(</span><span class="n">eot</span><span class="p">,</span> <span class="n">workers</span><span class="p">,</span> <span class="n">pool</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This thread is for monitoring subprocesses forked by GeneratorDataset/map/batch</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">workers</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;[Internal Error] The 2rd parameter of watch dog thread should be list of process, &quot;</span>\
                        <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">workers</span><span class="p">)))</span>
    <span class="k">if</span> <span class="n">pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">pool</span><span class="p">,</span> <span class="n">multiprocessing</span><span class="o">.</span><span class="n">pool</span><span class="o">.</span><span class="n">Pool</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;[Internal Error] The 3rd parameter of watch dog thread should be multiprocessing.Pool, &quot;</span>\
                        <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">pool</span><span class="p">)))</span>
    <span class="k">while</span> <span class="ow">not</span> <span class="n">eot</span><span class="o">.</span><span class="n">is_set</span><span class="p">():</span>
        <span class="n">subprocess_exit_num</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="c1"># Monitoring and count how many subprocesses already exit</span>
        <span class="n">subprocess_exit_num</span> <span class="o">=</span> <span class="n">_monitor_subprocess_exit</span><span class="p">(</span><span class="n">workers</span><span class="p">)</span>
        <span class="c1"># If find subprocess exit, we will wait for 30s and do some waitpid operations</span>
        <span class="k">if</span> <span class="n">subprocess_exit_num</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># Python multiprocessing.pool has a bug, if sub process of pool is killed, pool will launch</span>
                <span class="c1"># a new sub process, so we have to set worker_handler._state to TERMINATE to stop relaunching.</span>
                <span class="k">if</span> <span class="n">pool</span><span class="o">.</span><span class="n">_state</span> <span class="o">==</span> <span class="n">RUN</span><span class="p">:</span>  <span class="c1"># pylint: disable=W0212</span>
                    <span class="n">pool</span><span class="o">.</span><span class="n">_state</span> <span class="o">=</span> <span class="n">TERMINATE</span>  <span class="c1"># pylint: disable=W0212</span>
                    <span class="n">pool</span><span class="o">.</span><span class="n">_worker_handler</span><span class="o">.</span><span class="n">_state</span> <span class="o">=</span> <span class="n">TERMINATE</span>  <span class="c1"># pylint: disable=W0212</span>
            <span class="n">start</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span>
            <span class="k">while</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span> <span class="o">-</span> <span class="n">start</span> <span class="o">&lt;</span> <span class="mi">30</span><span class="p">:</span>
                <span class="c1"># We need to distinguishing get_dataset_size or train finished normally and hang scenario.</span>
                <span class="c1"># If get_dataset_size or train finished normally, _stop_subprocess can be execute and</span>
                <span class="c1"># self.need_abort can be set to True. If main process is hang in get(), self.need_abort</span>
                <span class="c1"># will never set to True, then we wait for 30s and kill main process</span>
                <span class="k">if</span> <span class="n">eot</span><span class="o">.</span><span class="n">is_set</span><span class="p">():</span>
                    <span class="k">return</span>
                <span class="c1"># Sometimes subprocess may be zombie, so in 30s we can wait and do some useful tasks(waitpid).</span>
                <span class="n">wait_pid</span><span class="p">()</span>
            <span class="c1"># multiprocessing.queue may hang in .get() forever when put() process was killed.</span>
            <span class="c1"># We have to exit main process otherwise main process will hang.</span>
            <span class="k">if</span> <span class="n">pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">_terminate_process</span><span class="p">(</span><span class="n">pool</span><span class="o">.</span><span class="n">_pool</span><span class="p">)</span>  <span class="c1"># pylint: disable=W0212</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">_terminate_process</span><span class="p">(</span><span class="n">workers</span><span class="p">)</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">critical</span><span class="p">(</span><span class="s2">&quot;The subprocess of dataset may exit unexpected or be killed, &quot;</span>
                            <span class="s2">&quot;main process will exit.&quot;</span><span class="p">)</span>
            <span class="n">os</span><span class="o">.</span><span class="n">kill</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">getpid</span><span class="p">(),</span> <span class="n">signal</span><span class="o">.</span><span class="n">SIGTERM</span><span class="p">)</span>


<span class="c1"># Pyfunc collection for multiprocess pyfunc</span>
<span class="c1"># This global variable will only be used within subprocesses</span>
<span class="n">_GLOBAL_PYFUNC_LIST</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">_ARGS_QUEUE</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">_RET_QUEUE</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">_OP_NAME</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
<span class="n">_OP_PROCESS</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
<span class="n">_LOCK</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Lock</span><span class="p">()</span>


<span class="c1"># Pyfunc worker init function</span>
<span class="c1"># Python multiprocessing library forbid sending lambda function through pipe.</span>
<span class="c1"># This init function allow us to add all Python function to a global collection and then fork afterwards.</span>
<span class="k">def</span> <span class="nf">_pyfunc_worker_init</span><span class="p">(</span><span class="n">pyfunc_list</span><span class="p">,</span> <span class="n">args_queue</span><span class="p">,</span> <span class="n">ret_queue</span><span class="p">):</span>
    <span class="k">global</span> <span class="n">_GLOBAL_PYFUNC_LIST</span>
    <span class="k">global</span> <span class="n">_ARGS_QUEUE</span>
    <span class="k">global</span> <span class="n">_RET_QUEUE</span>
    <span class="n">_GLOBAL_PYFUNC_LIST</span> <span class="o">=</span> <span class="n">pyfunc_list</span>
    <span class="n">_ARGS_QUEUE</span> <span class="o">=</span> <span class="n">args_queue</span>
    <span class="n">_RET_QUEUE</span> <span class="o">=</span> <span class="n">ret_queue</span>


<span class="c1"># Pyfunc worker execution function</span>
<span class="c1"># All exceptions will be raised to main processes</span>
<span class="k">def</span> <span class="nf">_pyfunc_worker_exec</span><span class="p">(</span><span class="n">index</span><span class="p">,</span> <span class="n">qid</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Internal function for call certain pyfunc in Python process.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># Some threads in multiprocess.pool can&#39;t process sigint signal,</span>
    <span class="c1"># and will occur hang problem, so ctrl+c will pass to parent process.</span>
    <span class="n">signal</span><span class="o">.</span><span class="n">signal</span><span class="p">(</span><span class="n">signal</span><span class="o">.</span><span class="n">SIGINT</span><span class="p">,</span> <span class="n">signal</span><span class="o">.</span><span class="n">SIG_IGN</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">qid</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span><span class="p">:</span>
        <span class="c1"># Pass arguments through the Queue instead of directly to remote process</span>
        <span class="n">args</span> <span class="o">=</span> <span class="n">_ARGS_QUEUE</span><span class="p">[</span><span class="n">qid</span><span class="p">]</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">r</span> <span class="o">=</span> <span class="n">_GLOBAL_PYFUNC_LIST</span><span class="p">[</span><span class="n">index</span><span class="p">](</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">ExceptionHandler</span><span class="p">(</span><span class="n">where</span><span class="o">=</span><span class="s2">&quot;in map(or batch) worker and execute python function&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="n">_RET_QUEUE</span><span class="p">[</span><span class="n">qid</span><span class="p">]</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">r</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">_RET_QUEUE</span><span class="p">[</span><span class="n">qid</span><span class="p">]</span><span class="o">.</span><span class="n">put</span><span class="p">((</span><span class="n">r</span><span class="p">,))</span>
        <span class="k">return</span> <span class="p">[</span><span class="n">qid</span><span class="p">]</span>
    <span class="c1"># not using shared memory for passing arguments, call function directly</span>
    <span class="n">result</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">_GLOBAL_PYFUNC_LIST</span><span class="p">[</span><span class="n">index</span><span class="p">](</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
    <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">ExceptionHandler</span><span class="p">(</span><span class="n">where</span><span class="o">=</span><span class="s2">&quot;in map(or batch) worker and execute python function&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span>


<span class="c1"># PythonCallable wrapper for multiprocess pyfunc</span>
<span class="k">class</span> <span class="nc">_PythonCallable</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Internal Python function wrapper for multiprocessing pyfunc.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">py_callable</span><span class="p">,</span> <span class="n">idx</span><span class="p">,</span> <span class="n">pool</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">arg_q</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">res_q</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="c1"># Original Python callable from user.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">py_callable</span> <span class="o">=</span> <span class="n">py_callable</span>
        <span class="c1"># Process pool created for current iterator.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pool</span> <span class="o">=</span> <span class="n">pool</span>
        <span class="c1"># Python callable index for subprocess _GLOBAL_PYFUNC_LIST</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">idx</span> <span class="o">=</span> <span class="n">idx</span>

        <span class="k">if</span> <span class="n">pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">queuemap</span> <span class="o">=</span> <span class="p">{}</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">arg_q</span> <span class="o">=</span> <span class="n">arg_q</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">res_q</span> <span class="o">=</span> <span class="n">res_q</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">next_queue</span> <span class="o">=</span> <span class="mi">0</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pool_is_running</span><span class="p">()</span> <span class="ow">and</span> <span class="n">check_iterator_cleanup</span><span class="p">()</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>
            <span class="n">result</span><span class="p">,</span> <span class="n">qid</span><span class="p">,</span> <span class="n">ret</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_send</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">ret</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">result</span>

            <span class="c1"># todo this check might be wrong</span>
            <span class="k">while</span> <span class="n">check_iterator_cleanup</span><span class="p">()</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>
                <span class="k">try</span><span class="p">:</span>
                    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_receive</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="n">qid</span><span class="p">)</span>
                <span class="k">except</span> <span class="n">multiprocessing</span><span class="o">.</span><span class="n">TimeoutError</span><span class="p">:</span>
                    <span class="k">continue</span>
                <span class="k">except</span> <span class="ne">KeyboardInterrupt</span><span class="p">:</span>
                    <span class="n">_set_iterator_cleanup</span><span class="p">()</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="o">.</span><span class="n">join</span><span class="p">()</span>
                    <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Multiprocess MapOp worker receives KeyboardInterrupt.&quot;</span><span class="p">)</span>
            <span class="k">return</span> <span class="p">(</span><span class="kc">None</span><span class="p">,)</span>
        <span class="c1"># Invoke original Python callable in master process in case the pool is gone.</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">py_callable</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">to_json</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">py_callable</span><span class="o">.</span><span class="n">to_json</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_send</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        The map/batch operator will use multiprocessing-pool apply_async interface to execute python function</span>
<span class="sd">        in a sub process, apply_async will release GIL temporarily. For better performance, we use shared memory</span>
<span class="sd">        feature and pass shared queue instead of multiprocess args.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="n">qid</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">arg_q</span> <span class="o">!=</span> <span class="p">[]:</span>
            <span class="n">tid</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">get_ident</span><span class="p">()</span>
            <span class="c1"># Need to register each thread to use a different queue to send data to pool</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">tid</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">queuemap</span><span class="p">:</span>
                <span class="n">qid</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_queue</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">next_queue</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_queue</span> <span class="o">+</span> <span class="mi">1</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">queuemap</span><span class="p">[</span><span class="n">tid</span><span class="p">]</span> <span class="o">=</span> <span class="n">qid</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">qid</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">queuemap</span><span class="p">[</span><span class="n">tid</span><span class="p">]</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">arg_q</span><span class="p">[</span><span class="n">qid</span><span class="p">]</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">args</span><span class="p">)</span>

            <span class="c1"># This call will send the tensors along with Python callable index to the process pool.</span>
            <span class="c1"># Block, yield GIL. Current thread will reacquire GIL once result is returned.</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pool_is_running</span><span class="p">()</span> <span class="ow">and</span> <span class="n">check_iterator_cleanup</span><span class="p">()</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>
                <span class="n">result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="o">.</span><span class="n">apply_async</span><span class="p">(</span><span class="n">_pyfunc_worker_exec</span><span class="p">,</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">idx</span><span class="p">,</span> <span class="n">qid</span><span class="p">,</span> <span class="p">[]])</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">ret</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">py_callable</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="o">.</span><span class="n">apply_async</span><span class="p">(</span><span class="n">_pyfunc_worker_exec</span><span class="p">,</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">idx</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">result</span><span class="p">,</span> <span class="n">qid</span><span class="p">,</span> <span class="n">ret</span>

    <span class="k">def</span> <span class="nf">_receive</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">result</span><span class="p">,</span> <span class="n">qid</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        The map/batch operator will use multiprocessing-pool get interface to sync output data from a sub process,</span>
<span class="sd">        get interface will reacquire GIL. For better performance, we use shared memory feature and get data from</span>
<span class="sd">        shared queue directly.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">arg_q</span> <span class="o">!=</span> <span class="p">[]:</span>
            <span class="n">r</span> <span class="o">=</span> <span class="n">result</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="mi">30</span><span class="p">)</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="n">ExceptionHandler</span><span class="p">):</span>
                <span class="n">r</span><span class="o">.</span><span class="n">reraise</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">r</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">qid</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;In PyCallable, got results from wrong thread&quot;</span><span class="p">)</span>
            <span class="n">r</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">res_q</span><span class="p">[</span><span class="n">qid</span><span class="p">]</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>
            <span class="k">return</span> <span class="n">r</span>
        <span class="n">r</span> <span class="o">=</span> <span class="n">result</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="mi">30</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="n">ExceptionHandler</span><span class="p">):</span>
            <span class="n">r</span><span class="o">.</span><span class="n">reraise</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">r</span>

    <span class="k">def</span> <span class="nf">_pool_is_running</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="c1"># note here: the RUN state of python3.7 and python3.8 is different:</span>
        <span class="c1"># python3.7: RUN = 0</span>
        <span class="c1"># python3.8: RUN = &quot;RUN&quot;</span>
        <span class="c1"># so we use self.pool._state == RUN instead and we can&#39;t use _state == 0 any more.</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">pool</span><span class="o">.</span><span class="n">_state</span> <span class="o">==</span> <span class="n">RUN</span><span class="p">:</span>  <span class="c1"># pylint: disable=W0212</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="k">return</span> <span class="kc">False</span>


<span class="k">def</span> <span class="nf">_mp_pool_exit_preprocess</span><span class="p">():</span>
    <span class="k">if</span> <span class="n">check_iterator_cleanup</span><span class="p">()</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>
        <span class="c1"># Set the iterator_cleanup flag to True before exiting, and wait 3s for all apply_async</span>
        <span class="c1"># applied to the multiprocessing task to prevent multiprocessing from hang when exiting</span>
        <span class="n">_set_iterator_cleanup</span><span class="p">()</span>
        <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_ExceptHookHandler</span><span class="p">:</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">excepthook</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">__handler_exception</span>

    <span class="k">def</span> <span class="nf">__handler_exception</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ex_type</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">tb</span><span class="p">):</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">critical</span><span class="p">(</span><span class="s2">&quot;Uncaught exception: &quot;</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="p">(</span><span class="n">ex_type</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">tb</span><span class="p">))</span>
        <span class="n">_mp_pool_exit_preprocess</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">MapDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying the Map operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be mapped.</span>
<span class="sd">        operations (TensorOp): A function mapping a nested structure of tensors</span>
<span class="sd">            to another nested structure of tensor (default=None).</span>
<span class="sd">        input_columns (Union[str, list[str]]): List of names of the input columns</span>
<span class="sd">            (default=None, the operations will be applied on the first columns in the dataset).</span>
<span class="sd">            The size of the list should match the number of inputs of the first operator.</span>
<span class="sd">        output_columns (Union[str, list[str]], optional): List of names of the output columns.</span>
<span class="sd">            The size of the list should match the number of outputs of the last operator</span>
<span class="sd">            (default=None, output columns will be the input columns, i.e., the columns will</span>
<span class="sd">            be replaced).</span>
<span class="sd">        column_order (list[str], optional): Specifies the list of all the columns you need in the whole</span>
<span class="sd">            dataset. The parameter is required when len(input_column) != len(output_column). Caution: the list here</span>
<span class="sd">            is not just the columns specified in parameter input_columns and output_columns.</span>
<span class="sd">        num_parallel_workers (int, optional): Number of workers to process the dataset</span>
<span class="sd">            in parallel (default=None).</span>
<span class="sd">        python_multiprocessing (bool, optional): Parallelize Python operations with multiple worker process. This</span>
<span class="sd">            option could be beneficial if the Python operation is computational heavy (default=False).</span>
<span class="sd">        cache (DatasetCache, optional): Use tensor caching service to speed up dataset processing.</span>
<span class="sd">            (default=None, which means no cache is used).</span>
<span class="sd">        callbacks (DSCallback, list[DSCallback], optional): List of Dataset callbacks to be called (Default=None)</span>
<span class="sd">        max_rowsize(int, optional): Maximum size of row in MB that is used for shared memory allocation to copy</span>
<span class="sd">            data between processes.  This is only used if python_multiprocessing is set to True (default=16).</span>
<span class="sd">        offload (bool, optional): Flag to indicate whether offload is used (Default=None).</span>

<span class="sd">    Raises:</span>
<span class="sd">        ValueError: If len(input_columns) != len(output_columns) and column_order is not specified.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">operations</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">input_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">output_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">column_order</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">python_multiprocessing</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">callbacks</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="mi">16</span><span class="p">,</span>
                 <span class="n">offload</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">operations</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">operations</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">operations</span> <span class="o">=</span> <span class="n">py_transforms</span><span class="o">.</span><span class="n">Compose</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">operations</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">input_columns</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">output_columns</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">column_order</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">column_order</span><span class="p">,</span> <span class="p">[])</span>

        <span class="c1">#  If output_columns were not provided then use input_columns</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span> <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span> \
                <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span><span class="p">)</span> \
                <span class="ow">and</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">column_order</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;When length of input_columns and output_columns are not equal,&quot;</span>
                             <span class="s2">&quot; column_order must be specified.&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">python_multiprocessing</span> <span class="o">=</span> <span class="n">python_multiprocessing</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">hook</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">workers</span> <span class="o">=</span> <span class="p">[]</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">callbacks</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">callbacks</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span> <span class="o">=</span> <span class="n">max_rowsize</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">offload</span> <span class="o">=</span> <span class="n">offload</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">operations</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">op</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">operations</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">op</span> <span class="ow">and</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="s1">&#39;parse&#39;</span><span class="p">,</span> <span class="kc">None</span><span class="p">):</span>
                <span class="n">operations</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">op</span><span class="o">.</span><span class="n">parse</span><span class="p">())</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">operations</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">op</span><span class="p">)</span>

        <span class="n">callbacks</span> <span class="o">=</span> <span class="p">[</span><span class="n">cb</span><span class="o">.</span><span class="n">create_runtime_obj</span><span class="p">()</span> <span class="k">for</span> <span class="n">cb</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">callbacks</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">MapNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">operations</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_columns</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">column_order</span><span class="p">,</span>
                           <span class="n">callbacks</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span><span class="p">,</span> <span class="n">OffloadToManualOffloadMode</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">offload</span><span class="p">])</span>

    <span class="k">def</span> <span class="nf">__deepcopy__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">memodict</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">__safe_deepcopy__</span><span class="p">(</span><span class="n">memodict</span><span class="p">,</span> <span class="n">exclude</span><span class="o">=</span><span class="p">(</span><span class="s2">&quot;operations&quot;</span><span class="p">,</span> <span class="s2">&quot;callbacks&quot;</span><span class="p">,</span> <span class="s2">&quot;__transfer_dataset__&quot;</span><span class="p">))</span>

    <span class="c1"># Iterator bootstrap will be called on iterator construction.</span>
    <span class="c1"># A deep copy of Dataset object is created prior of iterator_bootstrap.</span>
    <span class="c1"># This method will create per iterator process pool and bind pyfunc execution to the pool.</span>
    <span class="k">def</span> <span class="nf">iterator_bootstrap</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Per iterator bootstrap callback.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">python_multiprocessing</span><span class="p">:</span>
            <span class="n">iter_specific_operations</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">callable_list</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">arg_q_list</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">res_q_list</span> <span class="o">=</span> <span class="p">[]</span>

            <span class="c1"># If user didn&#39;t specify num_parallel_workers, set it to default</span>
            <span class="n">num_parallel</span> <span class="o">=</span> <span class="n">get_num_parallel_workers</span><span class="p">()</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">num_parallel</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span>

            <span class="k">if</span> <span class="n">get_enable_shared_mem</span><span class="p">():</span>
                <span class="n">_check_shm_usage</span><span class="p">(</span><span class="n">num_parallel</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
                <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_parallel</span><span class="p">):</span>
                    <span class="n">arg_q_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_SharedQueue</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span><span class="p">))</span>
                    <span class="n">res_q_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_SharedQueue</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_rowsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">max_rowsize</span><span class="p">))</span>

            <span class="c1"># Pass #1, look for Python callables and build list</span>
            <span class="k">for</span> <span class="n">op</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">operations</span><span class="p">:</span>
                <span class="c1"># our c transforms is now callable and should not be run in Python multithreading</span>
                <span class="k">if</span> <span class="n">MapDataset</span><span class="o">.</span><span class="n">__operation_valid_for_multiprocessing</span><span class="p">(</span><span class="n">op</span><span class="p">):</span>
                    <span class="n">callable_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">op</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">callable_list</span><span class="p">:</span>
                <span class="c1"># Construct pool with the callable list</span>
                <span class="c1"># The callable list and _pyfunc_worker_init are used to pass lambda function in to subprocesses</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="o">=</span> <span class="n">multiprocessing</span><span class="o">.</span><span class="n">Pool</span><span class="p">(</span><span class="n">processes</span><span class="o">=</span><span class="n">num_parallel</span><span class="p">,</span>
                                                         <span class="n">initializer</span><span class="o">=</span><span class="n">_pyfunc_worker_init</span><span class="p">,</span>
                                                         <span class="n">initargs</span><span class="o">=</span><span class="p">(</span><span class="n">callable_list</span><span class="p">,</span> <span class="n">arg_q_list</span><span class="p">,</span> <span class="n">res_q_list</span><span class="p">))</span>

                <span class="c1"># Pass #2</span>
                <span class="n">idx</span> <span class="o">=</span> <span class="mi">0</span>
                <span class="k">global</span> <span class="n">_OP_NAME</span><span class="p">,</span> <span class="n">_OP_PROCESS</span><span class="p">,</span> <span class="n">_LOCK</span>
                <span class="n">op_id</span> <span class="o">=</span> <span class="n">_OP_NAME</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="p">)]</span>
                <span class="c1"># obtain process id from multiprocessing.pool</span>
                <span class="n">process_id</span> <span class="o">=</span> <span class="p">{</span><span class="n">op_id</span><span class="p">:</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">num_parallel_workers</span><span class="p">,</span> <span class="nb">set</span><span class="p">()]}</span>
                <span class="k">for</span> <span class="n">pool</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">_pool</span><span class="p">:</span>  <span class="c1"># pylint: disable=W0212</span>
                    <span class="n">process_id</span><span class="p">[</span><span class="n">op_id</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">pool</span><span class="o">.</span><span class="n">pid</span><span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">workers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">pool</span><span class="p">)</span>
                <span class="k">with</span> <span class="n">_LOCK</span><span class="p">:</span>
                    <span class="n">_OP_PROCESS</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">process_id</span><span class="p">)</span>
                <span class="k">for</span> <span class="n">op</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">operations</span><span class="p">:</span>
                    <span class="c1"># our c transforms is now callable and should not be run in Python multithreading</span>
                    <span class="k">if</span> <span class="n">MapDataset</span><span class="o">.</span><span class="n">__operation_valid_for_multiprocessing</span><span class="p">(</span><span class="n">op</span><span class="p">):</span>
                        <span class="c1"># Wrap Python callable into _PythonCallable</span>
                        <span class="n">iter_specific_operations</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_PythonCallable</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="n">idx</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="p">,</span>
                                                                        <span class="n">arg_q_list</span><span class="p">,</span> <span class="n">res_q_list</span><span class="p">))</span>
                        <span class="n">idx</span> <span class="o">+=</span> <span class="mi">1</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="c1"># CPP ops remain the same</span>
                        <span class="n">iter_specific_operations</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">op</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">operations</span> <span class="o">=</span> <span class="n">iter_specific_operations</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">hook</span> <span class="o">=</span> <span class="n">_ExceptHookHandler</span><span class="p">()</span>

                <span class="c1"># Map multiprocessing will launch a watch dog thread for monitoring sub processes</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_launch_watch_dog</span><span class="p">()</span>

                <span class="n">atexit</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="n">_mp_pool_exit_preprocess</span><span class="p">)</span>
                <span class="c1"># If Python version greater than 3.8, we need to close ThreadPool in atexit for unclean pool teardown.</span>
                <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">version_info</span> <span class="o">&gt;=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">8</span><span class="p">):</span>
                    <span class="n">atexit</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">close</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">__operation_valid_for_multiprocessing</span><span class="p">(</span><span class="n">op</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="n">op</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">str</span><span class="p">(</span><span class="n">op</span><span class="p">)</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s2">&quot;c_transform&quot;</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">_launch_watch_dog</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">platform</span><span class="o">.</span><span class="n">system</span><span class="p">()</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="o">!=</span> <span class="s1">&#39;windows&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Event</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Thread</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">_watch_dog</span><span class="p">,</span> <span class="n">args</span><span class="o">=</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">eot</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">workers</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span><span class="o">.</span><span class="n">daemon</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_abort_watchdog</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">eot</span><span class="o">.</span><span class="n">is_set</span><span class="p">():</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">eot</span><span class="o">.</span><span class="n">set</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;process_pool&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process_pool</span><span class="o">.</span><span class="n">join</span><span class="p">()</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;watch_dog&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">watch_dog</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;eot&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">eot</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_abort_watchdog</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">FilterDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying filter predicate to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be mapped.</span>
<span class="sd">        predicate (callable): Python callable which returns a boolean value. If False then filter the element.</span>
<span class="sd">        input_columns (Union[str, list[str]], optional): List of names of the input columns</span>
<span class="sd">        (default=None, the predicate will be applied to all columns in the dataset).</span>
<span class="sd">        num_parallel_workers (int, optional): Number of workers to process the dataset</span>
<span class="sd">            in parallel (default=None).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">predicate</span><span class="p">,</span> <span class="n">input_columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">,</span> <span class="n">num_parallel_workers</span><span class="o">=</span><span class="n">num_parallel_workers</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">predicate</span> <span class="o">=</span> <span class="k">lambda</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="nb">bool</span><span class="p">(</span><span class="n">predicate</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">input_columns</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">FilterNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">predicate</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_columns</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">RepeatDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Repeat operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be repeated.</span>
<span class="sd">        count (int): Number of times the dataset will be repeated (default=-1, repeat indefinitely).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">count</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">count</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">count</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">RepeatNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">count</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SkipDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Skip operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input dataset to have elements skipped.</span>
<span class="sd">        count (int): Number of elements to be skipped in the dataset.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">count</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">count</span> <span class="o">=</span> <span class="n">count</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">SkipNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">count</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">TakeDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Take operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to have elements taken from.</span>
<span class="sd">        count (int): Number of elements to be taken from the dataset.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">count</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">count</span> <span class="o">=</span> <span class="n">count</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">TakeNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">count</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">ZipDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Zip operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        datasets (tuple): A tuple of datasets to be zipped together.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dataset is not an instance of Dataset.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">datasets</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">datasets</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">ZipNode</span><span class="p">(</span><span class="n">children</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">is_sync</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">any</span><span class="p">([</span><span class="n">c</span><span class="o">.</span><span class="n">is_sync</span><span class="p">()</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">])</span>


<span class="k">class</span> <span class="nc">ConcatDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying concat dataset operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        datasets (list): A list of datasets to be concatenated together.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dataset is not an instance of Dataset.</span>
<span class="sd">        ValueError: If there is no samples in the one of the datasets.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">datasets</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">datasets</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">dataset</span> <span class="ow">in</span> <span class="n">datasets</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">Dataset</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Invalid dataset, expected Dataset object, but got </span><span class="si">%s</span><span class="s2">!&quot;</span> <span class="o">%</span> <span class="nb">type</span><span class="p">(</span><span class="n">dataset</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">datasets</span> <span class="o">=</span> <span class="n">datasets</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_sampler</span> <span class="o">=</span> <span class="n">samplers</span><span class="o">.</span><span class="n">SequentialSampler</span><span class="p">(</span><span class="n">num_samples</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">children_sizes_</span> <span class="o">=</span> <span class="p">[</span><span class="n">c</span><span class="o">.</span><span class="n">get_dataset_size</span><span class="p">()</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">]</span>
        <span class="n">child_index</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">children_sizes_</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">item</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;There are no samples in the dataset number </span><span class="si">%d</span><span class="s2">. Please make sure there are &quot;</span>
                                 <span class="s2">&quot;valid samples in the dataset.&quot;</span> <span class="o">%</span> <span class="n">child_index</span><span class="p">)</span>
            <span class="n">child_index</span> <span class="o">+=</span> <span class="mi">1</span>

        <span class="c1"># _children_flag_and_nums: A list of pair&lt;int ,int&gt;.The first element of pair is flag that characterizes</span>
        <span class="c1"># whether the dataset is mappable. The second element of pair is length of the dataset</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span> <span class="o">=</span> <span class="p">[]</span>

        <span class="c1"># _children_start_end_index_: A list of pair&lt;int ,int&gt;.The elements of pair are used to characterize</span>
        <span class="c1"># the valid position of the dataset corresponding to the subscript when sampling</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">index</span><span class="p">,</span> <span class="n">child</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">):</span>
            <span class="n">tem_list</span> <span class="o">=</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">]</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">tem_list</span><span class="p">)</span>
            <span class="n">dataset_len</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">children_sizes_</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>

            <span class="kn">from</span> <span class="nn">mindspore.dataset.engine.datasets_user_defined</span> <span class="kn">import</span> <span class="n">GeneratorDataset</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">child</span><span class="p">,</span> <span class="n">GeneratorDataset</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">child</span><span class="o">.</span><span class="n">source</span><span class="p">,</span> <span class="s2">&quot;__getitem__&quot;</span><span class="p">):</span>
                <span class="n">dataset_len</span> <span class="o">=</span> <span class="mi">0</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">children_sizes_</span><span class="p">[</span><span class="n">index</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>

            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">child</span><span class="p">,</span> <span class="n">MappableDataset</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">dataset_len</span><span class="p">))</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="mi">1</span><span class="p">,</span> <span class="n">dataset_len</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">ConcatNode</span><span class="p">(</span><span class="n">children</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_sampler</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">use_sampler</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sampler</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the distributedSampler to concat dataset</span>

<span class="sd">        Args:</span>
<span class="sd">            sampler (Sampler): The sampler to use for the current dataset.</span>
<span class="sd">                Currently supported: DistributedSampler.</span>

<span class="sd">        Raises:</span>
<span class="sd">            TypeError: If the sampler is not an instance of DistributedSampler</span>
<span class="sd">            ValueError: If the parameter shuffle of sampler is True</span>
<span class="sd">            ValueError: If the parameter NumSamples of sampler is not None.</span>
<span class="sd">            ValueError: If num_shards &lt;=0.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">sampler</span><span class="p">,</span> <span class="n">samplers</span><span class="o">.</span><span class="n">DistributedSampler</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;The parameter </span><span class="si">%s</span><span class="s2"> of concat must be DistributedSampler!&quot;</span> <span class="o">%</span> <span class="n">sampler</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">sampler</span><span class="o">.</span><span class="n">is_shuffled</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The parameter shuffle of DistributedSampler must be False!&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">sampler</span><span class="o">.</span><span class="n">num_shards</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The parameter num_shards of DistributedSampler must be positive int!&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">sampler</span><span class="o">.</span><span class="n">get_num_samples</span><span class="p">()</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The parameter num_samples of DistributedSampler is not support to be set!&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_size</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_sampler</span> <span class="o">=</span> <span class="n">sampler</span>
        <span class="n">cumulative_samples_nums</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">index</span><span class="p">,</span> <span class="n">child</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">child</span><span class="p">,</span> <span class="s1">&#39;sampler&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="n">child</span><span class="o">.</span><span class="n">sampler</span><span class="o">.</span><span class="n">get_num_samples</span><span class="p">()</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The parameter NumSamples of </span><span class="si">%s</span><span class="s2"> is not support to be set!&quot;</span> <span class="o">%</span> <span class="n">child</span><span class="p">)</span>

            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">child</span><span class="p">,</span> <span class="n">BatchDataset</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;The parameter </span><span class="si">%s</span><span class="s2"> of concat must not be BatchDataset!&quot;</span> <span class="o">%</span> <span class="n">child</span><span class="p">)</span>

            <span class="c1"># if child is mappable and the length is greater than 0</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">1</span><span class="p">]:</span>

                <span class="n">tem_value</span> <span class="o">=</span> <span class="n">cumulative_samples_nums</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span>

                <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">_children_flag_and_nums</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="o">&gt;=</span> <span class="n">sampler</span><span class="o">.</span><span class="n">num_shards</span><span class="p">:</span>
                    <span class="k">if</span> <span class="n">tem_value</span> <span class="o">&lt;</span> <span class="n">sampler</span><span class="o">.</span><span class="n">num_shards</span><span class="p">:</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">cumulative_samples_nums</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">tem_value</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">cumulative_samples_nums</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">_children_start_end_index_</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">tem_value</span> <span class="o">%</span> <span class="n">sampler</span><span class="o">.</span><span class="n">num_shards</span>

                <span class="n">tem_sampler</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">sampler</span><span class="p">)</span>
                <span class="n">tem_sampler</span><span class="o">.</span><span class="n">set_offset</span><span class="p">(</span><span class="n">cumulative_samples_nums</span><span class="p">)</span>
                <span class="n">child</span><span class="o">.</span><span class="n">use_sampler</span><span class="p">(</span><span class="n">tem_sampler</span><span class="p">)</span>

            <span class="n">cumulative_samples_nums</span> <span class="o">+=</span> <span class="bp">self</span><span class="o">.</span><span class="n">children_sizes_</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
            <span class="n">cumulative_samples_nums</span> <span class="o">%=</span> <span class="n">sampler</span><span class="o">.</span><span class="n">num_shards</span>


<span class="k">class</span> <span class="nc">RenameDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Rename operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be Renamed.</span>
<span class="sd">        input_columns (Union[str, list[str]]): List of names of the input columns.</span>
<span class="sd">        output_columns (Union[str, list[str]]): List of names of the output columns.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">input_columns</span><span class="p">,</span> <span class="n">output_columns</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_column_names</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">input_columns</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_column_names</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">output_columns</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">RenameNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_column_names</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_column_names</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">to_list</span><span class="p">(</span><span class="n">items</span><span class="p">):</span>
    <span class="k">if</span> <span class="n">items</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">return</span> <span class="p">[]</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">items</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">list</span><span class="p">(</span><span class="n">items</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">items</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
        <span class="k">return</span> <span class="p">[</span><span class="n">items</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">items</span>


<span class="k">class</span> <span class="nc">ProjectDataset</span><span class="p">(</span><span class="n">UnionBaseDataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying Project operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be Projected.</span>
<span class="sd">        columns (Union[str, list[str]]): List of names of the columns to project.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">columns</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">columns</span> <span class="o">=</span> <span class="n">to_list</span><span class="p">(</span><span class="n">columns</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">ProjectNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">columns</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_ToDevice</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Internal class to handle sending data to device.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">num_epochs</span><span class="p">):</span>
        <span class="n">ir_tree</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">api_tree</span> <span class="o">=</span> <span class="n">dataset</span><span class="o">.</span><span class="n">create_ir_tree</span><span class="p">()</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_runtime_context</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">PythonRuntimeContext</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_runtime_context</span><span class="o">.</span><span class="n">Init</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">ToDevice</span><span class="p">(</span><span class="n">num_epochs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">Init</span><span class="p">(</span><span class="n">ir_tree</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_runtime_context</span><span class="o">.</span><span class="n">AssignConsumer</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="p">)</span>

        <span class="n">ITERATORS_LIST</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">weakref</span><span class="o">.</span><span class="n">ref</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span>
        <span class="n">_unset_iterator_cleanup</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">send</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">Send</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">stop_send</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        send stop send signal to pipeline, it is used when end of sequence is sent at the epoch end.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">StopSend</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">continue_send</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        send continue send signal to pipeline, it is used when end of sequence is sent at the epoch end.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">ContinueSend</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">get_data_info</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get type and shape of current batch.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">GetDataInfo</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">release</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Manually terminate Device Queue instead of relying on out of scope destruction.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;_runtime_context&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">_runtime_context</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;_to_device&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_runtime_context</span><span class="o">.</span><span class="n">Terminate</span><span class="p">()</span>
                <span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span>
            <span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">_runtime_context</span>

    <span class="k">def</span> <span class="nf">__deepcopy__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">memodict</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="nf">get_offload_model</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get offload model containing removed offload ops from pipeline.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">offload_model</span> <span class="o">=</span> <span class="n">GetOffloadModel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">offload_model</span>


<span class="k">class</span> <span class="nc">TransferDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The result of applying TDT operator to the input Dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        input_dataset (Dataset): Input Dataset to be transferred.</span>
<span class="sd">        send_epoch_end (bool, optional): Whether to send end of sequence to device or not (default=True).</span>
<span class="sd">        create_data_info_queue (bool, optional): Whether to create queue which stores</span>
<span class="sd">            types and shapes of data or not (default=False).</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If device_type is empty.</span>
<span class="sd">        ValueError: If device_type is not &#39;Ascend&#39;, &#39;GPU&#39; or &#39;CPU&#39;.</span>
<span class="sd">        RuntimeError: If dataset is unknown.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dataset</span><span class="p">,</span> <span class="n">send_epoch_end</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">create_data_info_queue</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">children</span><span class="o">=</span><span class="n">input_dataset</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">queue_name</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">uuid</span><span class="o">.</span><span class="n">uuid1</span><span class="p">())</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">device_type</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;device_target&quot;</span><span class="p">)</span> <span class="k">if</span> <span class="n">context</span> <span class="k">else</span> <span class="s2">&quot;CPU&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">device_id</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;device_id&quot;</span><span class="p">)</span> <span class="k">if</span> <span class="n">context</span> <span class="k">else</span> <span class="mi">0</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_send_epoch_end</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">send_epoch_end</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_create_data_info_queue</span> <span class="o">=</span> <span class="n">create_data_info_queue</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">total_batch</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s2">&quot;__total_batch__&quot;</span><span class="p">):</span>
            <span class="n">total_batch</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">__total_batch__</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">TransferNode</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">queue_name</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">device_type</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">device_id</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_send_epoch_end</span><span class="p">,</span>
                                <span class="n">total_batch</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_create_data_info_queue</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">create_dict_iterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_epochs</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">output_numpy</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;TransferDataset is not iterable.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">create_tuple_iterator</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_epochs</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">output_numpy</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">do_copy</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;TransferDataset is not iterable.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__iter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;TransferDataset is not iterable.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">output_shapes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;TransferDataset does not support obtaining output_shapes.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">output_types</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;TransferDataset does not support obtaining output_types.&quot;</span><span class="p">)</span>

    <span class="nd">@check_to_device_send</span>
    <span class="k">def</span> <span class="nf">send</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_epochs</span><span class="o">=-</span><span class="mi">1</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Send to device</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">Dataset</span><span class="o">.</span><span class="n">_noop_mode</span><span class="p">():</span>
            <span class="k">return</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="o">=</span> <span class="n">_ToDevice</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_epochs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">send</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">stop_send</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">stop_send</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">continue_send</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">continue_send</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">get_data_info</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get type and shape of current batch</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">get_data_info</span><span class="p">()</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Calling get_data_info with bad state.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">get_offload_model</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">get_offload_model</span><span class="p">()</span>

        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;get_offload_model, _to_device is None&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">release</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Manually terminate Device Queue instead of relying on out of scope destruction.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_to_device</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">Schema</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Class to represent a schema of a dataset.</span>

<span class="sd">    Args:</span>
<span class="sd">        schema_file(str): Path of the schema file (default=None).</span>

<span class="sd">    Returns:</span>
<span class="sd">        Schema object, schema info about dataset.</span>

<span class="sd">    Raises:</span>
<span class="sd">        RuntimeError: If schema file failed to load.</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Create schema; specify column name, mindspore.dtype and shape of the column</span>
<span class="sd">        &gt;&gt;&gt; schema = ds.Schema()</span>
<span class="sd">        &gt;&gt;&gt; schema.add_column(name=&#39;col1&#39;, de_type=mstype.int64, shape=[2])</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@check_schema</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">schema_file</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">schema_file</span> <span class="o">=</span> <span class="n">replace_none</span><span class="p">(</span><span class="n">schema_file</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cpp_schema</span> <span class="o">=</span> <span class="n">cde</span><span class="o">.</span><span class="n">SchemaObj</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">schema_file</span><span class="p">)</span>

    <span class="nd">@check_add_column</span>
    <span class="k">def</span> <span class="nf">add_column</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">de_type</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Add new column to the schema.</span>

<span class="sd">        Args:</span>
<span class="sd">            name (str): The new name of the column.</span>
<span class="sd">            de_type (str): Data type of the column.</span>
<span class="sd">            shape (list[int], optional): Shape of the column</span>
<span class="sd">                (default=None, [-1] which is an unknown shape of rank 1).</span>

<span class="sd">        Raises:</span>
<span class="sd">            ValueError: If column type is unknown.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">de_type</span><span class="p">,</span> <span class="n">typing</span><span class="o">.</span><span class="n">Type</span><span class="p">):</span>
            <span class="n">de_type</span> <span class="o">=</span> <span class="n">mstype_to_detype</span><span class="p">(</span><span class="n">de_type</span><span class="p">)</span>
            <span class="n">col_type</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">de_type</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">col_type</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">cde</span><span class="o">.</span><span class="n">DataType</span><span class="p">(</span><span class="n">de_type</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">shape</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cpp_schema</span><span class="o">.</span><span class="n">add_column</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">col_type</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cpp_schema</span><span class="o">.</span><span class="n">add_column</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">col_type</span><span class="p">,</span> <span class="n">shape</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">parse_columns</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Parse the columns and add it to self.</span>

<span class="sd">        Args:</span>
<span class="sd">            columns (Union[dict, list[dict], tuple[dict]]): Dataset attribute information, decoded from schema file.</span>

<span class="sd">                - list[dict], &#39;name&#39; and &#39;type&#39; must be in keys, &#39;shape&#39; optional.</span>

<span class="sd">                - dict, columns.keys() as name, columns.values() is dict, and &#39;type&#39; inside, &#39;shape&#39; optional.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: If failed to parse columns.</span>
<span class="sd">            RuntimeError: If column&#39;s name field is missing.</span>
<span class="sd">            RuntimeError: If column&#39;s type field is missing.</span>

<span class="sd">        Examples:</span>
<span class="sd">            &gt;&gt;&gt; from mindspore.dataset import Schema</span>
<span class="sd">            &gt;&gt;&gt; schema = Schema()</span>
<span class="sd">            &gt;&gt;&gt; columns1 = [{&#39;name&#39;: &#39;image&#39;, &#39;type&#39;: &#39;int8&#39;, &#39;shape&#39;: [3, 3]},</span>
<span class="sd">            ...             {&#39;name&#39;: &#39;label&#39;, &#39;type&#39;: &#39;int8&#39;, &#39;shape&#39;: [1]}]</span>
<span class="sd">            &gt;&gt;&gt; schema.parse_columns(columns1)</span>
<span class="sd">            &gt;&gt;&gt; columns2 = {&#39;image&#39;: {&#39;shape&#39;: [3, 3], &#39;type&#39;: &#39;int8&#39;}, &#39;label&#39;: {&#39;shape&#39;: [1], &#39;type&#39;: &#39;int8&#39;}}</span>
<span class="sd">            &gt;&gt;&gt; schema.parse_columns(columns2)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cpp_schema</span><span class="o">.</span><span class="n">parse_columns</span><span class="p">(</span><span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">columns</span><span class="p">,</span> <span class="n">indent</span><span class="o">=</span><span class="mi">2</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">to_json</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get a JSON string of the schema.</span>

<span class="sd">        Returns:</span>
<span class="sd">            str, JSON string of the schema.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">cpp_schema</span><span class="o">.</span><span class="n">to_json</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">from_json</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">json_obj</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get schema file from JSON object.</span>

<span class="sd">        Args:</span>
<span class="sd">            json_obj(dictionary): Object of JSON parsed.</span>

<span class="sd">        Raises:</span>
<span class="sd">            RuntimeError: if there is unknown item in the object.</span>
<span class="sd">            RuntimeError: if dataset type is missing in the object.</span>
<span class="sd">            RuntimeError: if columns are missing in the object.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cpp_schema</span><span class="o">.</span><span class="n">from_string</span><span class="p">(</span><span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">json_obj</span><span class="p">,</span> <span class="n">indent</span><span class="o">=</span><span class="mi">2</span><span class="p">))</span>

    <span class="k">def</span> <span class="fm">__str__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">to_json</span><span class="p">()</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">get_num_rows</span><span class="p">(</span><span class="n">schema</span><span class="p">):</span>
        <span class="n">schema_obj</span> <span class="o">=</span> <span class="n">schema</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">schema_obj</span><span class="p">,</span> <span class="n">Schema</span><span class="p">):</span>
            <span class="n">schema_obj</span> <span class="o">=</span> <span class="n">Schema</span><span class="p">(</span><span class="n">schema_obj</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">schema_obj</span><span class="o">.</span><span class="n">cpp_schema</span><span class="o">.</span><span class="n">get_num_rows</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">DeserializedDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_obj</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_obj</span> <span class="o">=</span> <span class="n">input_obj</span>

    <span class="k">def</span> <span class="nf">parse</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">children</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_obj</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="n">json_str</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_obj</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">Dataset</span><span class="o">.</span><span class="n">from_json_string</span><span class="p">(</span><span class="n">json_str</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">cde</span><span class="o">.</span><span class="n">Dataset</span><span class="o">.</span><span class="n">from_json_file</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_obj</span><span class="p">)</span>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>